aristanetworks · gmuloc · Dec 5, 2024 · Nov 25, 2024 · Nov 28, 2024 · Nov 29, 2024
diff --git a/anta/decorators.py b/anta/decorators.py
@@ -17,7 +17,8 @@
 F = TypeVar("F", bound=Callable[..., Any])
 
 
-def deprecated_test(new_tests: list[str] | None = None) -> Callable[[F], F]:
+# TODO: Remove this decorator in ANTA v2.0.0 in favor of deprecated_test_class
+def deprecated_test(new_tests: list[str] | None = None) -> Callable[[F], F]:  # pragma: no cover
     """Return a decorator to log a message of WARNING severity when a test is deprecated.
 
     Parameters
@@ -62,6 +63,52 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any:
     return decorator
 
 
+def deprecated_test_class(new_tests: list[str] | None = None) -> Callable[[type[AntaTest]], type[AntaTest]]:
+    """Return a decorator to log a message of WARNING severity when a test is deprecated.
+
+    Parameters
+    ----------
+    new_tests
+        A list of new test classes that should replace the deprecated test.
+
+    Returns
+    -------
+    Callable[[type], type]
+        A decorator that can be used to wrap test functions.
+
+    """
+
+    def decorator(cls: type[AntaTest]) -> type[AntaTest]:
+        """Actual decorator that logs the message.
+
+        Parameters
+        ----------
+        cls
+            The cls to be decorated.
+
+        Returns
+        -------
+        cls
+            The decorated cls.
+        """
+        orig_init = cls.__init__
+
+        def new_init(*args: Any, **kwargs: Any) -> None:
+            """Overload __init__ to generate a warning message for deprecation."""
+            if new_tests:
+                new_test_names = ", ".join(new_tests)
+                logger.warning("%s test is deprecated. Consider using the following new tests: %s.", cls.name, new_test_names)
+            else:
+                logger.warning("%s test is deprecated.", cls.name)
+            orig_init(*args, **kwargs)
+
+        # NOTE: we are ignoring mypy warning as we want to assign to a method here
+        cls.__init__ = new_init  # type: ignore[method-assign]
+        return cls
+
+    return decorator
+
+
 def skip_on_platforms(platforms: list[str]) -> Callable[[F], F]:
     """Return a decorator to skip a test based on the device's hardware model.
 

diff --git a/anta/input_models/stun.py b/anta/input_models/stun.py
@@ -0,0 +1,35 @@
+# 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.
+"""Module containing input models for services tests."""
+
+from __future__ import annotations
+
+from ipaddress import IPv4Address
+
+from pydantic import BaseModel, ConfigDict
+
+from anta.custom_types import Port
+
+
+class StunClientTranslation(BaseModel):
+    """STUN (Session Traversal Utilities for NAT) model represents the configuration of an IPv4-based client translations."""
+
+    model_config = ConfigDict(extra="forbid")
+    source_address: IPv4Address
+    """The IPv4 address of the STUN client"""
+    source_port: Port = 4500
+    """The port number used by the STUN client for communication. Defaults to 4500."""
+    public_address: IPv4Address | None = None
+    """The public-facing IPv4 address of the STUN client, discovered via the STUN server."""
+    public_port: Port | None = None
+    """The public-facing port number of the STUN client, discovered via the STUN server."""
+
+    def __str__(self) -> str:
+        """Return a human-readable string representation of the StunClientTranslation for reporting.
+
+        Examples
+        --------
+        Client 10.0.0.1 Port: 4500
+        """
+        return f"Client {self.source_address} Port: {self.source_port}"
diff --git a/anta/tests/stun.py b/anta/tests/stun.py
@@ -7,29 +7,36 @@
 # mypy: disable-error-code=attr-defined
 from __future__ import annotations
 
-from ipaddress import IPv4Address
 from typing import ClassVar
 
-from pydantic import BaseModel
-
-from anta.custom_types import Port
+from anta.decorators import deprecated_test_class
+from anta.input_models.stun import StunClientTranslation
 from anta.models import AntaCommand, AntaTemplate, AntaTest
-from anta.tools import get_failed_logs, get_value
+from anta.tools import get_value
+
 
+class VerifyStunClientTranslation(AntaTest):
+    """Verifies the translation for a source address on a STUN client.
 
-class VerifyStunClient(AntaTest):
-    """Verifies STUN client settings, including local IP/port and optionally public IP/port.
+    This test performs the following checks for each specified address family:
+
+      1. Validates that there is a translation for the source address on the STUN client.
+      2. If public IP and port details are provided, validates their correctness against the configuration.
 
     Expected Results
     ----------------
-    * Success: The test will pass if the STUN client is correctly configured with the specified IPv4 source address/port and public address/port.
-    * Failure: The test will fail if the STUN client is not configured or if the IPv4 source address, public address, or port details are incorrect.
+    * Success: If all of the following conditions are met:
+        - The test will pass if the source address translation is present.
+        - If public IP and port details are provided, they must also match the translation information.
+    * Failure: If any of the following occur:
+        - There is no translation for the source address on the STUN client.
+        - The public IP or port details, if specified, are incorrect.
 
     Examples
     --------
     ```yaml
     anta.tests.stun:
-      - VerifyStunClient:
+      - VerifyStunClientTranslation:
           stun_clients:
             - source_address: 172.18.3.2
               public_address: 172.18.3.21
@@ -43,73 +50,72 @@ class VerifyStunClient(AntaTest):
     """
 
     categories: ClassVar[list[str]] = ["stun"]
-    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaTemplate(template="show stun client translations {source_address} {source_port}")]
+    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaTemplate(template="show stun client translations {source_address} {source_port}", revision=1)]
 
     class Input(AntaTest.Input):
-        """Input model for the VerifyStunClient test."""
+        """Input model for the VerifyStunClientTranslation test."""
 
-        stun_clients: list[ClientAddress]
-
-        class ClientAddress(BaseModel):
-            """Source and public address/port details of STUN client."""
-
-            source_address: IPv4Address
-            """IPv4 source address of STUN client."""
-            source_port: Port = 4500
-            """Source port number for STUN client."""
-            public_address: IPv4Address | None = None
-            """Optional IPv4 public address of STUN client."""
-            public_port: Port | None = None
-            """Optional public port number for STUN client."""
+        stun_clients: list[StunClientTranslation]
+        """List of STUN clients."""
+        StunClientTranslation: ClassVar[type[StunClientTranslation]] = StunClientTranslation
 
     def render(self, template: AntaTemplate) -> list[AntaCommand]:
         """Render the template for each STUN translation."""
         return [template.render(source_address=client.source_address, source_port=client.source_port) for client in self.inputs.stun_clients]
 
     @AntaTest.anta_test
     def test(self) -> None:
-        """Main test function for VerifyStunClient."""
+        """Main test function for VerifyStunClientTranslation."""
         self.result.is_success()
 
         # Iterate over each command output and corresponding client input
         for command, client_input in zip(self.instance_commands, self.inputs.stun_clients):
             bindings = command.json_output["bindings"]
-            source_address = str(command.params.source_address)
-            source_port = command.params.source_port
+            input_public_address = client_input.public_address
+            input_public_port = client_input.public_port
 
             # If no bindings are found for the STUN client, mark the test as a failure and continue with the next client
             if not bindings:
-                self.result.is_failure(f"STUN client transaction for source `{source_address}:{source_port}` is not found.")
+                self.result.is_failure(f"{client_input} - STUN client translation not found.")
                 continue
 
-            # Extract the public address and port from the client input
-            public_address = client_input.public_address
-            public_port = client_input.public_port
-
             # Extract the transaction ID from the bindings
             transaction_id = next(iter(bindings.keys()))
 
-            # Prepare the actual and expected STUN data for comparison
-            actual_stun_data = {
-                "source ip": get_value(bindings, f"{transaction_id}.sourceAddress.ip"),
-                "source port": get_value(bindings, f"{transaction_id}.sourceAddress.port"),
-            }
-            expected_stun_data = {"source ip": source_address, "source port": source_port}
-
-            # If public address is provided, add it to the actual and expected STUN data
-            if public_address is not None:
-                actual_stun_data["public ip"] = get_value(bindings, f"{transaction_id}.publicAddress.ip")
-                expected_stun_data["public ip"] = str(public_address)
-
-            # If public port is provided, add it to the actual and expected STUN data
-            if public_port is not None:
-                actual_stun_data["public port"] = get_value(bindings, f"{transaction_id}.publicAddress.port")
-                expected_stun_data["public port"] = public_port
-
-            # If the actual STUN data does not match the expected STUN data, mark the test as failure
-            if actual_stun_data != expected_stun_data:
-    
2E18
            failed_log = get_failed_logs(expected_stun_data, actual_stun_data)
-                self.result.is_failure(f"For STUN source `{source_address}:{source_port}`:{failed_log}")
+            # Verifying the public address if provided
+            if input_public_address and str(input_public_address) != (actual_public_address := get_value(bindings, f"{transaction_id}.publicAddress.ip")):
+                self.result.is_failure(f"{client_input} - Incorrect public-facing address - Expected: {input_public_address} Actual: {actual_public_address}")
+
+            # Verifying the public port if provided
+            if input_public_port and input_public_port != (actual_public_port := get_value(bindings, f"{transaction_id}.publicAddress.port")):
+                self.result.is_failure(f"{client_input} - Incorrect public-facing port - Expected: {input_public_port} Actual: {actual_public_port}")
+
+
+@deprecated_test_class(new_tests=["VerifyStunClientTranslation"])
+class VerifyStunClient(VerifyStunClientTranslation):
+    """(Deprecated) Verifies the translation for a source address on a STUN client.
+
+    Alias for the VerifyStunClientTranslation test to maintain backward compatibility.
+    When initialized, it will emit a deprecation warning and call the VerifyStunClientTranslation test.
+
+    TODO: Remove this class in ANTA v2.0.0.
+
+    Examples
+    --------
+    ```yaml
+    anta.tests.stun:
+      - VerifyStunClient:
+          stun_clients:
+            - source_address: 172.18.3.2
+              public_address: 172.18.3.21
+              source_port: 4500
+              public_port: 6006
+    ```
+    """
+
+    # required to redefine name an description to overwrite parent class.
+    name = "VerifyStunClient"
+    description = "(Deprecated) Verifies the translation for a source address on a STUN client."
 
 
 class VerifyStunServer(AntaTest):

diff --git a/docs/api/tests.stun.md b/docs/api/tests.stun.md
@@ -7,6 +7,8 @@ anta_title: ANTA catalog for STUN tests
   ~ that can be found in the LICENSE file.
   -->
 
+# Tests
+
 ::: anta.tests.stun
     options:
       show_root_heading: false
@@ -18,3 +20,18 @@ anta_title: ANTA catalog for STUN tests
       filters:
         - "!test"
         - "!render"
+
+# Input models
+
+::: anta.input_models.stun
+
+    options:
+      show_root_heading: false
+      show_root_toc_entry: false
+      show_bases: false
+      merge_init_into_class: false
+      anta_hide_test_module_description: true
+      show_labels: true
+      filters:
+        - "!^__init__"
+        - "!^__str__"
diff --git a/examples/tests.yaml b/examples/tests.yaml
@@ -745,7 +745,14 @@ anta.tests.stp:
       threshold: 10
 anta.tests.stun:
   - VerifyStunClient:
-      # Verifies STUN client settings, including local IP/port and optionally public IP/port.
+      # (Deprecated) Verifies the translation for a source address on a STUN client.
+      stun_clients:
+        - source_address: 172.18.3.2
+          public_address: 172.18.3.21
+          source_port: 4500
+          public_port: 6006
+  - VerifyStunClientTranslation:
+      # Verifies the translation for a source address on a STUN client.
       stun_clients:
         - source_address: 172.18.3.2
           public_address: 172.18.3.21