add collected property to AntaCommand

anta/cli/debug/commands.py

@@ -34,24 +34,27 @@ def get_device(ctx: click.Context, param: Option, value: str) -> List[str]:
 
 @click.command()
 @click.option("--command", "-c", type=str, required=True, help="Command to run")
-@click.option("--ofmt", type=click.Choice(["text", "json"]), default="json", help="EOS eAPI format to use. can be text or json")
+@click.option("--ofmt", type=click.Choice(["json", "text"]), default="json", help="EOS eAPI format to use. can be text or json")
 @click.option("--api-version", "--version", type=EapiVersion(), default="latest", help="EOS eAPI version to use")
 @click.option("--device", "-d", type=str, required=True, help="Device from inventory to use", callback=get_device)
-def run_cmd(command: str, ofmt: str, api_version: Union[int, Literal["latest"]], device: AntaDevice) -> None:
+def run_cmd(command: str, ofmt: Literal["json", "text"], api_version: Union[int, Literal["latest"]], device: AntaDevice) -> None:
     """Run arbitrary command to an ANTA device"""
     console.print(f"Run command [green]{command}[/green] on [red]{device.name}[/red]")
     c = AntaCommand(command=command, ofmt=ofmt, version=api_version)
     asyncio.run(device.collect(c))
-    console.print(c.output)
+    if ofmt == 'json':
+        console.print(c.json_output)
+    if ofmt == 'text':
+        console.print(c.text_output)
 
 
 @click.command()
 @click.option("--template", "-t", type=str, required=True, help="Command template to run. E.g. 'show vlan {vlan_id}'")
-@click.option("--ofmt", type=click.Choice(["text", "json"]), default="json", help="EOS eAPI format to use. can be text or json")
+@click.option("--ofmt", type=click.Choice(["json", "text"]), default="json", help="EOS eAPI format to use. can be text or json")
 @click.option("--api-version", "--version", type=EapiVersion(), default="latest", help="EOS eAPI version to use")
 @click.option("--device", "-d", type=str, required=True, help="Device from inventory to use", callback=get_device)
 @click.argument("params", required=True, nargs=-1)
-def run_template(template: str, params: List[str], ofmt: str, api_version: Union[int, Literal["latest"]], device: AntaDevice) -> None:
+def run_template(template: str, params: List[str], ofmt: Literal["json", "text"], api_version: Union[int, Literal["latest"]], device: AntaDevice) -> None:
     """Run arbitrary templated command to an ANTA device.
 
     Takes a list of arguments (keys followed by a value) to build a dictionary used as template parameters.
@@ -62,8 +65,10 @@ def run_template(template: str, params: List[str], ofmt: str, api_version: Union
     template_params = dict(zip(params[::2], params[1::2]))
 
     console.print(f"Run templated command [blue]'{template}'[/blue] with [orange]{template_params}[/orange] on [red]{device.name}[/red]")
-    c = AntaCommand(
-        command=template.format(**template_params), template=AntaTemplate(template=template), template_params=template_params, ofmt=ofmt, version=api_version
-    )
+    t = AntaTemplate(template=template, params=template_params, ofmt=ofmt, version=api_version)
+    c = t.render(template_params)
     asyncio.run(device.collect(c))
-    console.print(c.output)
+    if ofmt == 'json':
+        console.print(c.json_output)
+    if ofmt == 'text':
+        console.print(c.text_output)

anta/cli/exec/utils.py

@@ -36,8 +36,8 @@ async def clear(dev: AntaDevice) -> None:
             commands.append(AntaCommand(command="clear hardware counter drop"))
         await dev.collect_commands(commands=commands)
         for command in commands:
-            if command.output is None:  # TODO - add a failed attribute to AntaCommand class
-                logger.error(f"Could not clear counters on device {dev.name}")
+            if not command.collected:
+                logger.error(f"Could not clear counters on device {dev.name}: {command.failed}")
         logger.info(f"Cleared counters on {dev.name} ({dev.hw_model})")
 
     logger.info("Connecting to devices...")

anta/decorators.py

@@ -88,7 +88,7 @@ async def wrapper(*args: Any, **kwargs: Dict[str, Any]) -> TestResult:
             await anta_test.device.collect(command=command)
 
             if command.failed is not None:
-                anta_test.result.is_error(f'{command.command}: {exc_to_str(command.failed)}')
+                anta_test.result.is_error(f"{command.command}: {exc_to_str(command.failed)}")
                 return anta_test.result
             if "vrfs" not in command.json_output:
                 anta_test.result.is_skipped(f"no BGP configuration for {family} on this device")

anta/device.py

@@ -3,7 +3,6 @@
 """
 import asyncio
 import logging
-import httpx
 from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Any, Dict, Iterator, List, Literal, Optional, Tuple, Union
@@ -201,6 +200,7 @@ def __init__(  # pylint: disable=R0913
             name = f"{host}:{port}"
         super().__init__(name, tags)
         self._enable_password = enable_password
+        timeout = httpx.Timeout(10.0, read=0.5)
         self._session: Device = Device(host=host, port=port, username=username, password=password, proto=proto, timeout=timeout)
         ssh_params: Dict[str, Any] = {}
         if insecure:
@@ -286,7 +286,7 @@ async def refresh(self) -> None:
         - hw_model: The hardware model of the device
         """
         # Refresh command
-        COMMAND: str = 'show version'
+        COMMAND: str = "show version"
         # Hardware model definition in show version
         HW_MODEL_KEY: str = "modelName"
         logger.debug(f"Refreshing device {self.name}")

anta/models.py

@@ -8,7 +8,7 @@
 from abc import ABC, abstractmethod
 from copy import deepcopy
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Coroutine, Dict, Literal, Optional, TypeVar, Union, List
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Coroutine, Dict, List, Literal, Optional, TypeVar, Union
 
 from pydantic import BaseModel
 
@@ -31,26 +31,32 @@ class AntaTemplate(BaseModel):
         template: Python f-string. Example: 'show vlan {vlan_id}'
         version: eAPI version - valid values are integers or the string "latest" - default is "latest"
         ofmt: eAPI output - json or text - default is json
-        vars: dictionary of variables with string values to render the Python f-string
+        params: dictionary of variables with string values to render the Python f-string
     """
 
     template: str
-    version: Union[int, Literal['latest']] = 'latest'
-    ofmt: Literal['json', 'text'] = 'json'
-    vars: Optional[Dict[str, str]]
-
-    def render(self, vars: Optional[Dict[str, str]] = None) -> AntaCommand:
-        if vars is None:
-            if self.vars is None:
-                raise RuntimeError(f'Cannot render template {self.template}: vars is missing')
+    version: Union[int, Literal["latest"]] = "latest"
+    ofmt: Literal["json", "text"] = "json"
+    params: Optional[Dict[str, str]]
+
+    def render(self, params: Optional[Dict[str, str]] = None) -> AntaCommand:
+        """Render an AntaCommand from an AntaTemplate instance.
+           Keep the parameters used in the AntaTemplate instance.
+
+            Args:
+                params: the template parameters. If not provided, will try to use the instance params if defined.
+
+            Returns:
+                AntaCommand: The rendered AntaCommand.
+                             This AntaCommand instance have a template attribute that references this
+                             AntaTemplate instance.
+        """
+        if params is None:
+            if self.params is None:
+                raise RuntimeError(f"Cannot render template {self.template}: params is missing")
         else:
-            self.vars = vars
-        return AntaCommand(
-                    command=self.template.format(**self.vars),
-                    ofmt=self.ofmt,
-                    version=self.version,
-                    template=self
-                )
+            self.params = params
+        return AntaCommand(command=self.template.format(**self.params), ofmt=self.ofmt, version=self.version, template=self)
 
 
 class AntaCommand(BaseModel):
@@ -60,41 +66,48 @@ class AntaCommand(BaseModel):
         command: Device command
         version: eAPI version - valid values are integers or the string "latest" - default is "latest"
         ofmt: eAPI output - json or text - default is json
-        output: collected output either dict for json or str for text
         template: AntaTemplate object used to render this command
         failed: If the command execution fails, the Exception object is stored in this field
     """
+
     class Config:
         # This is required if we want to keep an Exception object in the failed field
         arbitrary_types_allowed = True
 
     command: str
-    version: Union[int, Literal['latest']] = 'latest'
-    ofmt: Literal['json', 'text'] = 'json'
+    version: Union[int, Literal["latest"]] = "latest"
+    ofmt: Literal["json", "text"] = "json"
     output: Optional[Union[Dict[str, Any], str]]
     template: Optional[AntaTemplate]
     failed: Optional[Exception] = None
 
     @property
     def json_output(self) -> Dict[str, Any]:
+        """Get the command output as JSON"""
         if self.output is None:
-            raise RuntimeError(f'There is no output for command {self.command}')
-        if self.ofmt != 'json':
-            raise RuntimeError(f'Output of command {self.command} is not a JSON')
+            raise RuntimeError(f"There is no output for command {self.command}")
+        if self.ofmt != "json":
+            raise RuntimeError(f"Output of command {self.command} is not a JSON")
         if isinstance(self.output, str):
-            raise RuntimeError(f'Output of command {self.command} is invalid')
+            raise RuntimeError(f"Output of command {self.command} is invalid")
         return self.output
 
     @property
     def text_output(self) -> str:
+        """Get the command output as a string"""
         if self.output is None:
-            raise RuntimeError(f'There is no output for command {self.command}')
-        if self.ofmt != 'text':
-            raise RuntimeError(f'Output of command {self.command} is not a JSON')
+            raise RuntimeError(f"There is no output for command {self.command}")
+        if self.ofmt != "text":
+            raise RuntimeError(f"Output of command {self.command} is not a JSON")
         if not isinstance(self.output, str):
-            raise RuntimeError(f'Output of command {self.command} is invalid')
+            raise RuntimeError(f"Output of command {self.command} is invalid")
         return self.output
 
+    @property
+    def collected(self) -> bool:
+        """Return True if the command has been collected"""
+        return self.output is not None and self.failed is not None
+
 
 class AntaTestFilter(ABC):
     """Class to define a test Filter"""
@@ -180,7 +193,7 @@ def save_commands_data(self, eos_data: list[dict[Any, Any] | str]) -> None:
 
     def all_data_collected(self) -> bool:
         """returns True if output is populated for every command"""
-        return all(command.output is not None for command in self.instance_commands)
+        return all(command.collected for command in self.instance_commands)
 
     def get_failed_commands(self) -> List[AntaCommand]:
         """returns a list of all the commands that have a populated failed field"""
@@ -227,14 +240,14 @@ async def wrapper(
             **kwargs: dict[str, Any],
         ) -> TestResult:
             """
-                Wraps the test function and implement (in this order):
-                1. Instantiate the command outputs if `eos_data` is provided
-                2. Collect missing command outputs from the device
-                3. Run the test function
-                4. Catches and set the result if the test function raises an exception
-
-                Returns:
-                    TestResult: self.result, populated with the correct exit status
+            Wraps the test function and implement (in this order):
+            1. Instantiate the command outputs if `eos_data` is provided
+            2. Collect missing command outputs from the device
+            3. Run the test function
+            4. Catches and set the result if the test function raises an exception
+
+            Returns:
+                TestResult: self.result, populated with the correct exit status
             """
             if self.result.result != "unset":
                 return self.result
@@ -254,7 +267,7 @@ async def wrapper(
 
             try:
                 if cmds := self.get_failed_commands():
-                    self.result.is_error('\n'.join([f'{cmd.command}: {exc_to_str(cmd.failed)}' for cmd in cmds]))
+                    self.result.is_error("\n".join([f"{cmd.command}: {exc_to_str(cmd.failed)}" if cmd.failed else f"{cmd.command}: has failed" for cmd in cmds]))
                     return self.result
                 logger.debug(f"Test {self.name} on device {self.device.name}: running test")
                 function(self, **kwargs)