From 8853a026cbf0104f7726bbaa23d0c98cd1fcfd5a Mon Sep 17 00:00:00 2001 From: Carl Baillargeon Date: Tue, 6 Aug 2024 19:34:55 -0400 Subject: [PATCH] Added custom_test_specs --- python-avd/pyavd/_anta/utils/models.py | 6 ++++-- .../{anta_test_loader.py => test_loader.py} | 10 +++++----- .../utils/{anta_test_specs.py => test_specs.py} | 2 +- python-avd/pyavd/get_device_anta_catalog.py | 16 ++++++++++------ 4 files changed, 20 insertions(+), 14 deletions(-) rename python-avd/pyavd/_anta/utils/{anta_test_loader.py => test_loader.py} (90%) rename python-avd/pyavd/_anta/utils/{anta_test_specs.py => test_specs.py} (98%) diff --git a/python-avd/pyavd/_anta/utils/models.py b/python-avd/pyavd/_anta/utils/models.py index e6cfcbe05b9..c1a79f8888d 100644 --- a/python-avd/pyavd/_anta/utils/models.py +++ b/python-avd/pyavd/_anta/utils/models.py @@ -36,7 +36,8 @@ class FabricData: Attributes: ---------- structured_configs : dict[str, dict] - The structured configurations of the devices in the fabric. The key is the device name and the value is the structured config. + The structured configurations of the devices in the fabric. + The key is the device name and the value is the structured config. loopback0_mapping : dict[str, IPv4Address] The mapping of the Loopback0 IP addresses for each device. vtep_mapping : dict[str, IPv4Address] @@ -128,7 +129,8 @@ def check_inputs(self) -> Self: def is_ready(self) -> bool: """Check if the TestSpec has all necessary components to create an AntaTestDefinition. - A TestSpec is ready if the test class is defined and either an input factory or an input dict is provided if the test requires input. + A TestSpec is ready if the test class is defined and either an input factory or + an input dict is provided if the test requires input. """ is_ready = False diff --git a/python-avd/pyavd/_anta/utils/anta_test_loader.py b/python-avd/pyavd/_anta/utils/test_loader.py similarity index 90% rename from python-avd/pyavd/_anta/utils/anta_test_loader.py rename to python-avd/pyavd/_anta/utils/test_loader.py index 708c1388f84..5248adf3204 100644 --- a/python-avd/pyavd/_anta/utils/anta_test_loader.py +++ b/python-avd/pyavd/_anta/utils/test_loader.py @@ -9,12 +9,12 @@ from pyavd._utils import load_classes -from .anta_test_specs import ANTA_TEST_SPECS +from .test_specs import PYAVD_TEST_SPECS if TYPE_CHECKING: from anta.models import AntaTest - from .anta_test_specs import TestSpec + from .test_specs import TestSpec LOGGER = logging.getLogger("pyavd") @@ -44,9 +44,9 @@ def update_test_spec(test_spec: TestSpec, anta_available_tests: dict[str, type[A return True -ANTA_TEST_SPECS = list( +PYAVD_TEST_SPECS = list( filter( partial(update_test_spec, anta_available_tests=ANTA_AVAILABLE_TESTS), - ANTA_TEST_SPECS, - ) + PYAVD_TEST_SPECS, + ), ) diff --git a/python-avd/pyavd/_anta/utils/anta_test_specs.py b/python-avd/pyavd/_anta/utils/test_specs.py similarity index 98% rename from python-avd/pyavd/_anta/utils/anta_test_specs.py rename to python-avd/pyavd/_anta/utils/test_specs.py index 172e2c82f85..bf3a930b200 100644 --- a/python-avd/pyavd/_anta/utils/anta_test_specs.py +++ b/python-avd/pyavd/_anta/utils/test_specs.py @@ -16,7 +16,7 @@ from .constants import StructuredConfigKey from .models import TestSpec -ANTA_TEST_SPECS: list[TestSpec] = [ +PYAVD_TEST_SPECS: list[TestSpec] = [ TestSpec( name="VerifyAPIHttpsSSL", conditional_keys=[StructuredConfigKey.HTTPS_SSL_PROFILE], diff --git a/python-avd/pyavd/get_device_anta_catalog.py b/python-avd/pyavd/get_device_anta_catalog.py index 5724faf5c2d..6ce9f569a25 100644 --- a/python-avd/pyavd/get_device_anta_catalog.py +++ b/python-avd/pyavd/get_device_anta_catalog.py @@ -8,10 +8,10 @@ if TYPE_CHECKING: from anta.catalog import AntaCatalog - from ._anta.utils import FabricData + from ._anta.utils import FabricData, TestSpec -def get_device_anta_catalog(hostname: str, fabric_data: FabricData) -> AntaCatalog: +def get_device_anta_catalog(hostname: str, fabric_data: FabricData, custom_test_specs: list[TestSpec] | None = None) -> AntaCatalog: """Generate an ANTA catalog for a single device. Parameters @@ -22,8 +22,8 @@ def get_device_anta_catalog(hostname: str, fabric_data: FabricData) -> AntaCatal Contains relevant data (e.g. structured configurations, loopback mappings, etc.) of all devices in the fabric to generate the catalog. The instance must be created using the `get_fabric_data` function of this module. - logger : logging.Logger - Optional custom logger to use. If not provided, the `pyavd` logger will be used. + custom_test_specs : list[TestSpec] + Optional user-defined list of TestSpec to be added to the default PyAVD test specs. Returns: ------- @@ -31,12 +31,16 @@ def get_device_anta_catalog(hostname: str, fabric_data: FabricData) -> AntaCatal The generated ANTA catalog for the device. """ from ._anta.utils import ConfigManager, create_catalog - from ._anta.utils.anta_test_loader import ANTA_TEST_SPECS + from ._anta.utils.test_loader import PYAVD_TEST_SPECS + + custom_test_specs = custom_test_specs or [] # Create the device-specific ConfigManager used to generate the inputs for the tests config_manager = ConfigManager(hostname, fabric_data) - return create_catalog(config_manager, ANTA_TEST_SPECS) + PYAVD_TEST_SPECS.extend([test for test in custom_test_specs if test not in PYAVD_TEST_SPECS]) + + return create_catalog(config_manager, PYAVD_TEST_SPECS) def get_fabric_data(structured_configs: dict[str, dict]) -> FabricData: