diff options
Diffstat (limited to 'anta/models.py')
| -rw-r--r-- | anta/models.py | 335 |
1 files changed, 214 insertions, 121 deletions
diff --git a/anta/models.py b/anta/models.py index c8acda3..f963dc0 100644 --- a/anta/models.py +++ b/anta/models.py @@ -1,9 +1,8 @@ # Copyright (c) 2023-2024 Arista Networks, Inc. # Use of this source code is governed by the Apache License 2.0 # that can be found in the LICENSE file. -""" -Models to define a TestStructure -""" +"""Models to define a TestStructure.""" + from __future__ import annotations import hashlib @@ -14,77 +13,99 @@ from abc import ABC, abstractmethod from copy import deepcopy from datetime import timedelta from functools import wraps +from string import Formatter +from typing import TYPE_CHECKING, Any, Callable, ClassVar, Literal, TypeVar -# Need to keep Dict and List for pydantic in python 3.8 -from typing import TYPE_CHECKING, Any, Callable, ClassVar, Coroutine, Dict, List, Literal, Optional, TypeVar, Union - -from pydantic import BaseModel, ConfigDict, ValidationError, conint -from rich.progress import Progress, TaskID +from pydantic import BaseModel, ConfigDict, ValidationError, create_model from anta import GITHUB_SUGGESTION -from anta.logger import anta_log_exception +from anta.custom_types import Revision +from anta.logger import anta_log_exception, exc_to_str from anta.result_manager.models import TestResult -from anta.tools.misc import exc_to_str if TYPE_CHECKING: + from collections.abc import Coroutine + + from rich.progress import Progress, TaskID + from anta.device import AntaDevice F = TypeVar("F", bound=Callable[..., Any]) # Proper way to type input class - revisit this later if we get any issue @gmuloc # This would imply overhead to define classes # https://stackoverflow.com/questions/74103528/type-hinting-an-instance-of-a-nested-class -# N = TypeVar("N", bound="AntaTest.Input") - -# TODO - make this configurable - with an env var maybe? +# TODO: make this configurable - with an env var maybe? BLACKLIST_REGEX = [r"^reload.*", r"^conf\w*\s*(terminal|session)*", r"^wr\w*\s*\w+"] logger = logging.getLogger(__name__) -class AntaMissingParamException(Exception): - """ - This Exception should be used when an expected key in an AntaCommand.params dictionary - was not found. +class AntaParamsBaseModel(BaseModel): + """Extends BaseModel and overwrite __getattr__ to return None on missing attribute.""" - This Exception should in general never be raised in normal usage of ANTA. - """ + model_config = ConfigDict(extra="forbid") - def __init__(self, message: str) -> None: - self.message = "\n".join([message, GITHUB_SUGGESTION]) - super().__init__(self.message) + if not TYPE_CHECKING: + # Following pydantic declaration and keeping __getattr__ only when TYPE_CHECKING is false. + # Disabling 1 Dynamically typed expressions (typing.Any) are disallowed in `__getattr__ + # ruff: noqa: ANN401 + def __getattr__(self, item: str) -> Any: + """For AntaParams if we try to access an attribute that is not present We want it to be None.""" + try: + return super().__getattr__(item) + except AttributeError: + return None class AntaTemplate(BaseModel): """Class to define a command template as Python f-string. + Can render a command from parameters. - Attributes: + Attributes + ---------- template: Python f-string. Example: 'show vlan {vlan_id}' - version: eAPI version - valid values are 1 or "latest" - default is "latest" + version: eAPI version - valid values are 1 or "latest". revision: Revision of the command. Valid values are 1 to 99. Revision has precedence over version. - ofmt: eAPI output - json or text - default is json - use_cache: Enable or disable caching for this AntaTemplate if the AntaDevice supports it - default is True + ofmt: eAPI output - json or text. + use_cache: Enable or disable caching for this AntaTemplate if the AntaDevice supports it. + """ template: str version: Literal[1, "latest"] = "latest" - revision: Optional[conint(ge=1, le=99)] = None # type: ignore + revision: Revision | None = None ofmt: Literal["json", "text"] = "json" use_cache: bool = True - def render(self, **params: dict[str, Any]) -> AntaCommand: + def render(self, **params: str | int | bool) -> AntaCommand: """Render an AntaCommand from an AntaTemplate instance. + Keep the parameters used in the AntaTemplate instance. Args: + ---- params: dictionary of variables with string values to render the Python f-string - Returns: + Returns + ------- command: The rendered AntaCommand. This AntaCommand instance have a template attribute that references this AntaTemplate instance. + """ + # Create params schema on the fly + field_names = [fname for _, fname, _, _ in Formatter().parse(self.template) if fname] + # Extracting the type from the params based on the expected field_names from the template + fields: dict[str, Any] = {key: (type(params.get(key)), ...) for key in field_names} + # Accepting ParamsSchema as non lowercase variable + ParamsSchema = create_model( # noqa: N806 + "ParamsSchema", + __base__=AntaParamsBaseModel, + **fields, + ) + try: return AntaCommand( command=self.template.format(**params), @@ -92,7 +113,7 @@ class AntaTemplate(BaseModel): version=self.version, revision=self.revision, template=self, - params=params, + params=ParamsSchema(**params), use_cache=self.use_cache, ) except KeyError as e: @@ -113,70 +134,115 @@ class AntaCommand(BaseModel): __Revision has precedence over version.__ - Attributes: + Attributes + ---------- command: Device command - version: eAPI version - valid values are 1 or "latest" - default is "latest" + version: eAPI version - valid values are 1 or "latest". revision: eAPI revision of the command. Valid values are 1 to 99. Revision has precedence over version. - ofmt: eAPI output - json or text - default is json - output: Output of the command populated by the collect() function - template: AntaTemplate object used to render this command - params: Dictionary of variables with string values to render the template - errors: If the command execution fails, eAPI returns a list of strings detailing the error - use_cache: Enable or disable caching for this AntaCommand if the AntaDevice supports it - default is True + ofmt: eAPI output - json or text. + output: Output of the command. Only defined if there was no errors. + template: AntaTemplate object used to render this command. + errors: If the command execution fails, eAPI returns a list of strings detailing the error(s). + params: Pydantic Model containing the variables values used to render the template. + use_cache: Enable or disable caching for this AntaCommand if the AntaDevice supports it. + """ command: str version: Literal[1, "latest"] = "latest" - revision: Optional[conint(ge=1, le=99)] = None # type: ignore + revision: Revision | None = None ofmt: Literal["json", "text"] = "json" - output: Optional[Union[Dict[str, Any], str]] = None - template: Optional[AntaTemplate] = None - errors: List[str] = [] - params: Dict[str, Any] = {} + output: dict[str, Any] | str | None = None + template: AntaTemplate | None = None + errors: list[str] = [] + params: AntaParamsBaseModel = AntaParamsBaseModel() use_cache: bool = True @property def uid(self) -> str: - """Generate a unique identifier for this command""" + """Generate a unique identifier for this command.""" uid_str = f"{self.command}_{self.version}_{self.revision or 'NA'}_{self.ofmt}" - return hashlib.sha1(uid_str.encode()).hexdigest() + # Ignoring S324 probable use of insecure hash function - sha1 is enough for our needs. + return hashlib.sha1(uid_str.encode()).hexdigest() # noqa: S324 @property def json_output(self) -> dict[str, Any]: - """Get the command output as JSON""" + """Get the command output as JSON.""" if self.output is None: - raise RuntimeError(f"There is no output for command {self.command}") + msg = f"There is no output for command '{self.command}'" + raise RuntimeError(msg) if self.ofmt != "json" or not isinstance(self.output, dict): - raise RuntimeError(f"Output of command {self.command} is invalid") + msg = f"Output of command '{self.command}' is invalid" + raise RuntimeError(msg) return dict(self.output) @property def text_output(self) -> str: - """Get the command output as a string""" + """Get the command output as a string.""" if self.output is None: - raise RuntimeError(f"There is no output for command {self.command}") + msg = f"There is no output for command '{self.command}'" + raise RuntimeError(msg) if self.ofmt != "text" or not isinstance(self.output, str): - raise RuntimeError(f"Output of command {self.command} is invalid") + msg = f"Output of command '{self.command}' is invalid" + raise RuntimeError(msg) return str(self.output) @property + def error(self) -> bool: + """Return True if the command returned an error, False otherwise.""" + return len(self.errors) > 0 + + @property def collected(self) -> bool: - """Return True if the command has been collected""" - return self.output is not None and not self.errors + """Return True if the command has been collected, False otherwise. + + A command that has not been collected could have returned an error. + See error property. + """ + return not self.error and self.output is not None + + @property + def requires_privileges(self) -> bool: + """Return True if the command requires privileged mode, False otherwise. + + Raises + ------ + RuntimeError + If the command has not been collected and has not returned an error. + AntaDevice.collect() must be called before this property. + """ + if not self.collected and not self.error: + msg = f"Command '{self.command}' has not been collected and has not returned an error. Call AntaDevice.collect()." + raise RuntimeError(msg) + return any("privileged mode required" in e for e in self.errors) + + @property + def supported(self) -> bool: + """Return True if the command is supported on the device hardware platform, False otherwise. + + Raises + ------ + RuntimeError + If the command has not been collected and has not returned an error. + AntaDevice.collect() must be called before this property. + """ + if not self.collected and not self.error: + msg = f"Command '{self.command}' has not been collected and has not returned an error. Call AntaDevice.collect()." + raise RuntimeError(msg) + return not any("not supported on this hardware platform" in e for e in self.errors) class AntaTemplateRenderError(RuntimeError): - """ - Raised when an AntaTemplate object could not be rendered - because of missing parameters - """ + """Raised when an AntaTemplate object could not be rendered because of missing parameters.""" - def __init__(self, template: AntaTemplate, key: str): - """Constructor for AntaTemplateRenderError + def __init__(self, template: AntaTemplate, key: str) -> None: + """Initialize an AntaTemplateRenderError. Args: + ---- template: The AntaTemplate instance that failed to render key: Key that has not been provided to render the template + """ self.template = template self.key = key @@ -184,12 +250,13 @@ class AntaTemplateRenderError(RuntimeError): class AntaTest(ABC): - """Abstract class defining a test in ANTA + """Abstract class defining a test in ANTA. The goal of this class is to handle the heavy lifting and make writing a test as simple as possible. - Examples: + Examples + -------- The following is an example of an AntaTest subclass implementation: ```python class VerifyReachability(AntaTest): @@ -227,22 +294,24 @@ class AntaTest(ABC): instance_commands: List of AntaCommand instances of this test result: TestResult instance representing the result of this test logger: Python logger for this test instance + """ # Mandatory class attributes - # TODO - find a way to tell mypy these are mandatory for child classes - maybe Protocol + # TODO: find a way to tell mypy these are mandatory for child classes - maybe Protocol name: ClassVar[str] description: ClassVar[str] categories: ClassVar[list[str]] - commands: ClassVar[list[Union[AntaTemplate, AntaCommand]]] + commands: ClassVar[list[AntaTemplate | AntaCommand]] # Class attributes to handle the progress bar of ANTA CLI - progress: Optional[Progress] = None - nrfu_task: Optional[TaskID] = None + progress: Progress | None = None + nrfu_task: TaskID | None = None class Input(BaseModel): """Class defining inputs for a test in ANTA. - Examples: + Examples + -------- A valid test catalog will look like the following: ```yaml <Python module>: @@ -255,72 +324,85 @@ class AntaTest(ABC): ``` Attributes: result_overwrite: Define fields to overwrite in the TestResult object + """ model_config = ConfigDict(extra="forbid") - result_overwrite: Optional[ResultOverwrite] = None - filters: Optional[Filters] = None + result_overwrite: ResultOverwrite | None = None + filters: Filters | None = None def __hash__(self) -> int: - """ - Implement generic hashing for AntaTest.Input. + """Implement generic hashing for AntaTest.Input. + This will work in most cases but this does not consider 2 lists with different ordering as equal. """ return hash(self.model_dump_json()) class ResultOverwrite(BaseModel): - """Test inputs model to overwrite result fields + """Test inputs model to overwrite result fields. - Attributes: + Attributes + ---------- description: overwrite TestResult.description categories: overwrite TestResult.categories custom_field: a free string that will be included in the TestResult object + """ model_config = ConfigDict(extra="forbid") - description: Optional[str] = None - categories: Optional[List[str]] = None - custom_field: Optional[str] = None + description: str | None = None + categories: list[str] | None = None + custom_field: str | None = None class Filters(BaseModel): - """Runtime filters to map tests with list of tags or devices + """Runtime filters to map tests with list of tags or devices. + + Attributes + ---------- + tags: Tag of devices on which to run the test. - Attributes: - tags: List of device's tags for the test. """ model_config = ConfigDict(extra="forbid") - tags: Optional[List[str]] = None + tags: set[str] | None = None def __init__( self, device: AntaDevice, inputs: dict[str, Any] | AntaTest.Input | None = None, eos_data: list[dict[Any, Any] | str] | None = None, - ): - """AntaTest Constructor + ) -> None: + """AntaTest Constructor. Args: + ---- device: AntaDevice instance on which the test will be run inputs: dictionary of attributes used to instantiate the AntaTest.Input instance eos_data: Populate outputs of the test commands instead of collecting from devices. This list must have the same length and order than the `instance_commands` instance attribute. + """ self.logger: logging.Logger = logging.getLogger(f"{self.__module__}.{self.__class__.__name__}") self.device: AntaDevice = device self.inputs: AntaTest.Input self.instance_commands: list[AntaCommand] = [] - self.result: TestResult = TestResult(name=device.name, test=self.name, categories=self.categories, description=self.description) + self.result: TestResult = TestResult( + name=device.name, + test=self.name, + categories=self.categories, + description=self.description, + ) self._init_inputs(inputs) if self.result.result == "unset": self._init_commands(eos_data) def _init_inputs(self, inputs: dict[str, Any] | AntaTest.Input | None) -> None: - """Instantiate the `inputs` instance attribute with an `AntaTest.Input` instance - to validate test inputs from defined model. + """Instantiate the `inputs` instance attribute with an `AntaTest.Input` instance to validate test inputs using the model. + Overwrite result fields based on `ResultOverwrite` input definition. - Any input validation error will set this test result status as 'error'.""" + Any input validation error will set this test result status as 'error'. + """ try: if inputs is None: self.inputs = self.Input() @@ -340,10 +422,11 @@ class AntaTest(ABC): self.result.description = res_ow.description self.result.custom_field = res_ow.custom_field - def _init_commands(self, eos_data: Optional[list[dict[Any, Any] | str]]) -> None: + def _init_commands(self, eos_data: list[dict[Any, Any] | str] | None) -> None: """Instantiate the `instance_commands` instance attribute from the `commands` class attribute. + - Copy of the `AntaCommand` instances - - Render all `AntaTemplate` instances using the `render()` method + - Render all `AntaTemplate` instances using the `render()` method. Any template rendering error will set this test result status as 'error'. Any exception in user code in `render()` will set this test result status as 'error'. @@ -371,11 +454,11 @@ class AntaTest(ABC): return if eos_data is not None: - self.logger.debug(f"Test {self.name} initialized with input data") + self.logger.debug("Test %s initialized with input data", self.name) self.save_commands_data(eos_data) def save_commands_data(self, eos_data: list[dict[str, Any] | str]) -> None: - """Populate output of all AntaCommand instances in `instance_commands`""" + """Populate output of all AntaCommand instances in `instance_commands`.""" if len(eos_data) > len(self.instance_commands): self.result.is_error(message="Test initialization error: Trying to save more data than there are commands for the test") return @@ -386,11 +469,12 @@ class AntaTest(ABC): self.instance_commands[index].output = data def __init_subclass__(cls) -> None: - """Verify that the mandatory class attributes are defined""" + """Verify that the mandatory class attributes are defined.""" mandatory_attributes = ["name", "description", "categories", "commands"] for attr in mandatory_attributes: if not hasattr(cls, attr): - raise NotImplementedError(f"Class {cls.__module__}.{cls.__name__} is missing required class attribute {attr}") + msg = f"Class {cls.__module__}.{cls.__name__} is missing required class attribute {attr}" + raise NotImplementedError(msg) @property def collected(self) -> bool: @@ -400,15 +484,17 @@ class AntaTest(ABC): @property def failed_commands(self) -> list[AntaCommand]: """Returns a list of all the commands that have failed.""" - return [command for command in self.instance_commands if command.errors] + return [command for command in self.instance_commands if command.error] def render(self, template: AntaTemplate) -> list[AntaCommand]: - """Render an AntaTemplate instance of this AntaTest using the provided - AntaTest.Input instance at self.inputs. + """Render an AntaTemplate instance of this AntaTest using the provided AntaTest.Input instance at self.inputs. This is not an abstract method because it does not need to be implemented if there is - no AntaTemplate for this test.""" - raise NotImplementedError(f"AntaTemplate are provided but render() method has not been implemented for {self.__module__}.{self.name}") + no AntaTemplate for this test. + """ + _ = template + msg = f"AntaTemplate are provided but render() method has not been implemented for {self.__module__}.{self.name}" + raise NotImplementedError(msg) @property def blocked(self) -> bool: @@ -417,15 +503,17 @@ class AntaTest(ABC): for command in self.instance_commands: for pattern in BLACKLIST_REGEX: if re.match(pattern, command.command): - self.logger.error(f"Command <{command.command}> is blocked for security reason matching {BLACKLIST_REGEX}") + self.logger.error( + "Command <%s> is blocked for security reason matching %s", + command.command, + BLACKLIST_REGEX, + ) self.result.is_error(f"<{command.command}> is blocked for security reason") state = True return state async def collect(self) -> None: - """ - Method used to collect outputs of all commands of this test class from the device of this test instance. - """ + """Collect outputs of all commands of this test class from the device of this test instance.""" try: if self.blocked is False: await self.device.collect_commands(self.instance_commands) @@ -439,8 +527,7 @@ class AntaTest(ABC): @staticmethod def anta_test(function: F) -> Callable[..., Coroutine[Any, Any, TestResult]]: - """ - Decorator for the `test()` method. + """Decorate the `test()` method in child classes. This decorator implements (in this order): @@ -454,15 +541,21 @@ class AntaTest(ABC): async def wrapper( self: AntaTest, eos_data: list[dict[Any, Any] | str] | None = None, - **kwargs: Any, + **kwargs: dict[str, Any], ) -> TestResult: - """ + """Inner function for the anta_test decorator. + Args: + ---- + self: The test instance. eos_data: Populate outputs of the test commands instead of collecting from devices. This list must have the same length and order than the `instance_commands` instance attribute. + kwargs: Any keyword argument to pass to the test. - Returns: + Returns + ------- result: TestResult instance attribute populated with error status if any + """ def format_td(seconds: float, digits: int = 3) -> str: @@ -476,7 +569,7 @@ class AntaTest(ABC): # Data if eos_data is not None: self.save_commands_data(eos_data) - self.logger.debug(f"Test {self.name} initialized with input data {eos_data}") + self.logger.debug("Test %s initialized with input data %s", self.name, eos_data) # If some data is missing, try to collect if not self.collected: @@ -485,11 +578,10 @@ class AntaTest(ABC): return self.result if cmds := self.failed_commands: - self.logger.debug(self.device.supports) - unsupported_commands = [f"Skipped because {c.command} is not supported on {self.device.hw_model}" for c in cmds if not self.device.supports(c)] - self.logger.debug(unsupported_commands) + unsupported_commands = [f"'{c.command}' is not supported on {self.device.hw_model}" for c in cmds if not c.supported] if unsupported_commands: - self.logger.warning(f"Test {self.name} has been skipped because it is not supported on {self.device.hw_model}: {GITHUB_SUGGESTION}") + msg = f"Test {self.name} has been skipped because it is not supported on {self.device.hw_model}: {GITHUB_SUGGESTION}" + self.logger.warning(msg) self.result.is_skipped("\n".join(unsupported_commands)) return self.result self.result.is_error(message="\n".join([f"{c.command} has failed: {', '.join(c.errors)}" for c in cmds])) @@ -506,7 +598,8 @@ class AntaTest(ABC): self.result.is_error(message=exc_to_str(e)) test_duration = time.time() - start_time - self.logger.debug(f"Executing test {self.name} on device {self.device.name} took {format_td(test_duration)}") + msg = f"Executing test {self.name} on device {self.device.name} took {format_td(test_duration)}" + self.logger.debug(msg) AntaTest.update_progress() return self.result @@ -514,21 +607,20 @@ class AntaTest(ABC): return wrapper @classmethod - def update_progress(cls) -> None: - """ - Update progress bar for all AntaTest objects if it exists - """ + def update_progress(cls: type[AntaTest]) -> None: + """Update progress bar for all AntaTest objects if it exists.""" if cls.progress and (cls.nrfu_task is not None): cls.progress.update(cls.nrfu_task, advance=1) @abstractmethod def test(self) -> Coroutine[Any, Any, TestResult]: - """ - This abstract method is the core of the test logic. - It must set the correct status of the `result` instance attribute - with the appropriate outcome of the test. + """Core of the test logic. + + This is an abstractmethod that must be implemented by child classes. + It must set the correct status of the `result` instance attribute with the appropriate outcome of the test. - Examples: + Examples + -------- It must be implemented using the `AntaTest.anta_test` decorator: ```python @AntaTest.anta_test @@ -536,6 +628,7 @@ class AntaTest(ABC): self.result.is_success() for command in self.instance_commands: if not self._test_command(command): # _test_command() is an arbitrary test logic - self.result.is_failure("Failure reson") + self.result.is_failure("Failure reason") ``` + """ |