Consolidation of Genie parsing logic to Connection Broker

nac_test/pyats_core/broker/broker_client.py

@@ -126,6 +126,7 @@ async def _send_request(self, request: dict[str, Any]) -> dict[str, Any]:
 
             # Read response data
             response_data = await self.reader.readexactly(response_length)
+            logger.debug(f"Received response of length: {response_length} bytes")
             response = json.loads(response_data.decode("utf-8"))
 
             # Check for errors
@@ -165,6 +166,37 @@ async def execute_command(self, hostname: str, command: str) -> str:
 
         return result  # type: ignore[no-any-return]
 
+    async def parse_command(
+        self, hostname: str, command: str, return_raw: bool = False
+    ) -> dict[str, Any]:
+        """Parse command output on device through broker.
+
+        Args:
+            hostname: Device hostname
+            command: Command to parse
+            return_raw: Whether to return raw output along with parsed result
+
+        Returns:
+            Tuple containing parsed output dictionary and optionally raw output
+
+        Raises:
+            ConnectionError: If broker communication fails
+        """
+        logger.debug(
+            f"Parsing command on {hostname}: {command} (return_raw={return_raw})"
+        )
+
+        response = await self._send_request(
+            {
+                "command": "parse",
+                "hostname": hostname,
+                "cmd": command,
+                "return_raw": return_raw,
+            }
+        )
+
+        return response.get("result", {})  # type: ignore[no-any-return]
+
     async def ensure_connection(self, hostname: str) -> bool:
         """Ensure device is connected through broker.
 
@@ -258,6 +290,23 @@ async def execute(self, command: str) -> str:
         """
         return await self.broker_client.execute_command(self.hostname, command)
 
+    async def parse(
+        self, command: str, return_raw: bool = False
+    ) -> tuple[dict[str, Any], str | None]:
+        """Parse command output on device via broker.
+
+        Args:
+            command: Command to parse
+            return_raw: Whether to return raw output along with parsed result
+
+        Returns:
+            Dictionary containing parsed output and optionally raw output
+        """
+        result = await self.broker_client.parse_command(
+            self.hostname, command, return_raw
+        )
+        return result.get("parsed", {}), result.get("raw", None)
+
     async def connect(self) -> None:
         """Ensure device connection via broker."""
         success = await self.broker_client.ensure_connection(self.hostname)

nac_test/pyats_core/broker/connection_broker.py

@@ -199,6 +199,19 @@ async def _process_request(self, message: dict[str, Any]) -> dict[str, Any]:
 
                 result = await self._execute_command(hostname, cmd_string)
                 return {"status": "success", "result": result}
+            elif command == "parse":
+                hostname = message.get("hostname")
+                return_raw = message.get("return_raw", False)
+                cmd = message.get("cmd")
+
+                if not hostname or not cmd:
+                    return {
+                        "status": "error",
+                        "error": "Missing hostname or cmd parameter",
+                    }
+
+                result = await self._parse_command(hostname, cmd, return_raw)
+                return {"status": "success", "result": result}
 
             elif command == "connect":
                 hostname = message.get("hostname")
@@ -253,9 +266,10 @@ async def _execute_command(self, hostname: str, cmd: str) -> str:
         # Check cache first
         cached_output = cache.get(cmd)
         if cached_output is not None:
-            self.stats_command_cache_hits += 1
-            logger.debug(f"Broker cache hit for '{cmd}' on {hostname}")
-            return cached_output
+            if cached_output.output is not None:
+                self.stats_command_cache_hits += 1
+                logger.debug(f"Broker cache hit for '{cmd}' on {hostname}")
+                return cached_output.output
 
         # Command not in cache, need to execute
         self.stats_command_cache_misses += 1
@@ -271,7 +285,7 @@ async def _execute_command(self, hostname: str, cmd: str) -> str:
             output_str = str(output)
 
             # Cache the output for future requests
-            cache.set(cmd, output_str)
+            cache.set(cmd, output=output_str)
             logger.info(
                 f"Cached command output for '{cmd}' on {hostname} ({len(output_str)} chars)"
             )
@@ -283,6 +297,58 @@ async def _execute_command(self, hostname: str, cmd: str) -> str:
             await self._disconnect_device(hostname)
             raise
 
+    async def _parse_command(self, hostname: str, cmd: str, return_raw: bool) -> Any:
+        """Execute command and return parsed output using Genie parsers."""
+        # Get or create cache for this device
+        if hostname not in self.command_cache:
+            self.command_cache[hostname] = CommandCache(
+                hostname, ttl=3600
+            )  # 1 hour TTL
+            logger.info(f"Created command cache for device: {hostname}")
+
+        cache = self.command_cache[hostname]
+        cache_entry = cache.get(cmd)
+        if cache_entry is not None:
+            if cache_entry.parsed_output is not None:
+                self.stats_command_cache_hits += 1
+                logger.debug(f"Broker cache hit for parsed '{cmd}' on {hostname}")
+                return {
+                    "parsed": cache_entry.parsed_output,
+                    "raw": cache_entry.output if return_raw else None,
+                }
+
+        # Command not in cache, need to execute and parse
+        self.stats_command_cache_misses += 1
+        logger.debug(
+            f"Broker cache miss for parsed '{cmd}' on {hostname}, executing..."
+        )
+
+        # Ensure device is connected
+        connection = await self._get_connection(hostname)
+
+        # Execute command in thread pool (since Unicon is synchronous)
+        loop = get_or_create_event_loop()
+        if return_raw:
+            try:
+                output = await loop.run_in_executor(None, connection.execute, cmd)
+            except Exception as e:
+                logger.error(f"Command execution failed on {hostname}: {e}")
+                # Try to reconnect on failure
+                await self._disconnect_device(hostname)
+                raise
+        else:
+            output = None  # Don't retrieve raw output if not requested
+        try:
+            parsed_output = await loop.run_in_executor(None, connection.parse, cmd)
+            cache.set(cmd, output=output, parsed_output=parsed_output)
+            return {
+                "parsed": dict(parsed_output),
+                "raw": output if return_raw else None,
+            }
+        except Exception as e:
+            logger.error(f"Command parsing failed on {hostname}: {e}")
+            raise
+
     async def _get_connection(self, hostname: str) -> Any:
         """Get or create connection to device."""
         if hostname not in self.connection_locks:

nac_test/pyats_core/common/ssh_base_test.py

@@ -247,38 +247,28 @@ async def _async_setup(self, hostname: str) -> None:
         self.device_data = self.device_info
         # hostname already set in setup_ssh_context
 
-    def parse_output(
-        self, command: str, output: str | None = None
-    ) -> dict[str, Any] | None:
+    async def parse_output(
+        self, command: str, return_raw: bool = False
+    ) -> tuple[dict[str, Any] | None, str | None]:
         """Parse command output using Genie parser if available.
 
         This method attempts to use Genie parsers when a PyATS testbed is available.
         If no testbed is available or parsing fails, it returns None.
 
         Args:
             command: The command whose output should be parsed
-            output: Optional pre-fetched command output. If not provided,
-                   the command will be executed.
 
         Returns:
-            Parsed output dictionary if successful, None otherwise.
+            Tuple containing parsed output dictionary and optionally raw output if successful, None otherwise.
         """
-        # If we have a testbed device, use its parse method
-        if self.testbed_device:
-            try:
-                if output is not None:
-                    # Parse provided output
-                    result = self.testbed_device.parse(command, output=output)
-                    return dict(result) if result is not None else None
-                else:
-                    # Execute and parse in one step
-                    result = self.testbed_device.parse(command)
-                    return dict(result) if result is not None else None
-            except Exception as e:
-                self.logger.warning(f"Genie parser failed for '{command}': {e}")
-                return None
+
+        if self.connection is not None:
+            return await self.connection.parse(command, return_raw=return_raw)
+        elif self.testbed_device:
+            parsed_output = dict(self.testbed_device.parse(command))
+            return parsed_output, None
         else:
-            return None
+            return None, None
 
     def _create_execute_command_method(
         self, connection: Any, command_cache: CommandCache
@@ -307,10 +297,11 @@ async def execute_command(command: str) -> str:
             # Check cache first
             cached_output = command_cache.get(command)
             if cached_output is not None:
-                logging.debug(f"Using cached output for command: {command}")
-                # Track cached command execution for reporting
-                test_instance._track_ssh_command(command, cached_output)
-                return cached_output
+                if cached_output.output is not None:
+                    logging.debug(f"Using cached output for command: {command}")
+                    # Track cached command execution for reporting
+                    test_instance._track_ssh_command(command, cached_output.output)
+                    return cached_output.output
 
             # Execute command via connection (broker or testbed device)
             logging.debug(f"Executing command: {command}")
@@ -329,7 +320,7 @@ async def execute_command(command: str) -> str:
             output_str = str(output)
 
             # Cache the output
-            command_cache.set(command, output_str)
+            command_cache.set(command, output=output_str)
 
             # Track the command execution for reporting
             test_instance._track_ssh_command(command, output_str)

nac_test/pyats_core/ssh/command_cache.py

@@ -14,6 +14,15 @@
 logger = logging.getLogger(__name__)
 
 
+class CacheEntry:
+    """Represents a single cache entry for a command output."""
+
+    def __init__(self, output: str | None = None, parsed_output: Any | None = None):
+        self.output: str | None = output
+        self.parsed_output: Any | None = parsed_output
+        self.timestamp: float = time.time()
+
+
 class CommandCache:
     """Per-device command output cache with TTL support.
 
@@ -34,11 +43,11 @@ def __init__(self, hostname: str, ttl: int = 3600):
         """
         self.hostname = hostname
         self.ttl = ttl
-        self.cache: dict[str, dict[str, Any]] = {}  # command -> {output, timestamp}
+        self.cache: dict[str, CacheEntry] = {}  # command -> CacheEntry
 
         logger.debug(f"Initialized command cache for device {hostname} with TTL {ttl}s")
 
-    def get(self, command: str) -> str | None:
+    def get(self, command: str) -> CacheEntry | None:
         """Get cached command output if valid.
 
         Args:
@@ -49,27 +58,33 @@ def get(self, command: str) -> str | None:
         """
         if command in self.cache:
             entry = self.cache[command]
-            if time.time() - entry["timestamp"] < self.ttl:
+            if time.time() - entry.timestamp < self.ttl:
                 logger.debug(f"Cache hit for '{command}' on {self.hostname}")
-                return str(entry["output"])
+                return entry
             else:
                 # Entry has expired, remove it
                 del self.cache[command]
                 logger.debug(f"Cache expired for '{command}' on {self.hostname}")
 
         return None
 
-    def set(self, command: str, output: str) -> None:
+    def set(
+        self, command: str, output: str | None = None, parsed_output: Any | None = None
+    ) -> None:
         """Cache command output with current timestamp.
 
         Args:
             command: The command that was executed
             output: The command output to cache
+            parsed_output: The parsed command output to cache (optional)
         """
-        self.cache[command] = {"output": output, "timestamp": time.time()}
-        logger.debug(
-            f"Cached '{command}' output for {self.hostname} ({len(output)} chars)"
-        )
+        self.cache[command] = CacheEntry(output, parsed_output)
+        if output is not None:
+            logger.debug(
+                f"Cached '{command}' output for {self.hostname} ({len(output)} chars)"
+            )
+        if parsed_output is not None:
+            logger.debug(f"Cached '{command}' parsed output for {self.hostname}")
 
     def clear(self) -> None:
         """Clear all cached entries for this device."""
@@ -91,7 +106,7 @@ def get_cache_stats(self) -> dict[str, int]:
         valid_count = 0
 
         for entry in self.cache.values():
-            if current_time - entry["timestamp"] >= self.ttl:
+            if current_time - entry.timestamp >= self.ttl:
                 expired_count += 1
             else:
                 valid_count += 1

tests/e2e/fixtures/failure/templates/tests/verify_iosxe_control_fail.py

@@ -80,12 +80,8 @@ async def verify_item(self, semaphore, client, context):
                 start_time = time.time()
 
                 try:
-                    with self.test_context(api_context):
-                        output = await self.execute_command(command)
-                    command_duration = time.time() - start_time
-
                     parse_start = time.time()
-                    parsed_output = self.parse_output(command, output=output)
+                    parsed_output, _ = await self.parse_output(command)
                     parse_duration = time.time() - parse_start
 
                 except Exception as e:
@@ -101,7 +97,7 @@ async def verify_item(self, semaphore, client, context):
                         api_duration=api_duration,
                     )
 
-                api_duration = command_duration + parse_duration
+                api_duration = parse_duration
                 context["api_context"] = api_context
 
                 # Check if parsed output is empty or None

tests/e2e/fixtures/mixed/templates/tests/verify_iosxe_control.py

@@ -124,12 +124,9 @@ async def verify_item(self, semaphore, client, context):
                 start_time = time.time()
 
                 try:
-                    with self.test_context(api_context):
-                        output = await self.execute_command(command)
-                    command_duration = time.time() - start_time
 
                     parse_start = time.time()
-                    parsed_output = self.parse_output(command, output=output)
+                    parsed_output, _ = await self.parse_output(command)
                     parse_duration = time.time() - parse_start
 
                 except Exception as e:
@@ -155,7 +152,7 @@ async def verify_item(self, semaphore, client, context):
                         api_duration=api_duration,
                     )
 
-                api_duration = command_duration + parse_duration
+                api_duration = parse_duration
 
                 context["api_context"] = api_context
 

tests/e2e/fixtures/pyats_cc/templates/tests/verify_iosxe_no_critical_errors_in_system_logs.py

@@ -138,15 +138,11 @@ async def verify_item(self, semaphore, client, context):
 
                 start_time = time.time()
 
-                with self.test_context(api_context):
-                    output = await self.execute_command(command)
-                command_duration = time.time() - start_time
-
                 parse_start = time.time()
-                parsed_output = self.parse_output(command, output=output)
+                parsed_output, _ = await self.parse_output(command)
                 parse_duration = time.time() - parse_start
 
-                api_duration = command_duration + parse_duration
+                api_duration = parse_duration
 
                 context["api_context"] = api_context