[PATCH v4 06/11] dts: add NIC capability support

["Juraj Linkeš" <juraj.linkes@pantheon.tech>]

Some test cases or suites may be testing a NIC feature that is not
supported on all NICs, so add support for marking test cases or suites
as requiring NIC capabilities.

The marking is done with a decorator, which populates the internal
required_capabilities attribute of TestProtocol. The NIC capability
itself is a wrapper around the NicCapability defined in testpmd_shell.
The reason is Enums cannot be extended and the class implements the
methods of its abstract base superclass.

The decorator API is designed to be simple to use. The arguments passed
to it are all from the testpmd shell. Everything else (even the actual
capability object creation) is done internally.

Signed-off-by: Juraj Linkeš <juraj.linkes@pantheon.tech>
Reviewed-by: Dean Marx <dmarx@iol.unh.edu>
---
 dts/framework/remote_session/testpmd_shell.py |  55 +++++-
 dts/framework/testbed_model/capability.py     | 167 +++++++++++++++++-
 2 files changed, 220 insertions(+), 2 deletions(-)

diff --git a/dts/framework/remote_session/testpmd_shell.py b/dts/framework/remote_session/testpmd_shell.py
index 77902a468d..3401adcc28 100644
--- a/dts/framework/remote_session/testpmd_shell.py
+++ b/dts/framework/remote_session/testpmd_shell.py
@@ -17,10 +17,16 @@
 import functools
 import re
 import time
+from collections.abc import Callable, MutableSet
 from dataclasses import dataclass, field
 from enum import Flag, auto
 from pathlib import PurePath
-from typing import Any, Callable, ClassVar, Concatenate, ParamSpec
+from typing import TYPE_CHECKING, Any, ClassVar, Concatenate, ParamSpec, TypeAlias
+
+if TYPE_CHECKING:
+    from enum import Enum as NoAliasEnum
+else:
+    from aenum import NoAliasEnum
 
 from typing_extensions import Self, Unpack
 
@@ -37,6 +43,12 @@
 P = ParamSpec("P")
 TestPmdShellMethod = Callable[Concatenate["TestPmdShell", P], Any]
 
+TestPmdShellCapabilityMethod: TypeAlias = Callable[
+    ["TestPmdShell", MutableSet["NicCapability"], MutableSet["NicCapability"]], None
+]
+
+TestPmdShellDecorator: TypeAlias = Callable[[TestPmdShellMethod], TestPmdShellMethod]
+
 
 class TestPmdDevice:
     """The data of a device that testpmd can recognize.
@@ -986,3 +998,44 @@ def _close(self) -> None:
         self.stop()
         self.send_command("quit", "Bye...")
         return super()._close()
+
+
+class NicCapability(NoAliasEnum):
+    """A mapping between capability names and the associated :class:`TestPmdShell` methods.
+
+    The :class:`TestPmdShell` capability checking method executes the command that checks
+    whether the capability is supported.
+    A decorator may optionally be added to the method that will add and remove configuration
+    that's necessary to retrieve the capability support status.
+    The Enum members may be assigned the method itself or a tuple of
+    (capability_checking_method, decorator_function).
+
+    The signature of each :class:`TestPmdShell` capability checking method must be::
+
+        fn(self, supported_capabilities: MutableSet, unsupported_capabilities: MutableSet) -> None
+
+    The capability checking method must get whether a capability is supported or not
+    from a testpmd command. If multiple capabilities can be obtained from a testpmd command,
+    each should be obtained in the method. These capabilities should then
+    be added to `supported_capabilities` or `unsupported_capabilities` based on their support.
+
+    The two dictionaries are shared across all capability discovery function calls in a given
+    test run so that we don't call the same function multiple times.
+    """
+
+    def __call__(
+        self,
+        testpmd_shell: TestPmdShell,
+        supported_capabilities: MutableSet[Self],
+        unsupported_capabilities: MutableSet[Self],
+    ) -> None:
+        """Execute the associated capability retrieval function.
+
+        Args:
+            testpmd_shell: :class:`TestPmdShell` object to which the function will be bound to.
+            supported_capabilities: The dictionary storing the supported capabilities
+                of a given test run.
+            unsupported_capabilities: The dictionary storing the unsupported capabilities
+                of a given test run.
+        """
+        self.value(testpmd_shell, supported_capabilities, unsupported_capabilities)
diff --git a/dts/framework/testbed_model/capability.py b/dts/framework/testbed_model/capability.py
index 8899f07f76..fceec4440e 100644
--- a/dts/framework/testbed_model/capability.py
+++ b/dts/framework/testbed_model/capability.py
@@ -5,14 +5,29 @@
 
 This module provides a protocol that defines the common attributes of test cases and suites
 and support for test environment capabilities.
+
+Many test cases are testing features not available on all hardware.
+
+The module also allows developers to mark test cases or suites as requiring certain
+hardware capabilities with the :func:`requires` decorator.
 """
 
 from abc import ABC, abstractmethod
-from collections.abc import Sequence
+from collections.abc import MutableSet, Sequence
+from dataclasses import dataclass
 from typing import Callable, ClassVar, Protocol
 
 from typing_extensions import Self
 
+from framework.logger import get_dts_logger
+from framework.remote_session.testpmd_shell import (
+    NicCapability,
+    TestPmdShell,
+    TestPmdShellCapabilityMethod,
+    TestPmdShellDecorator,
+    TestPmdShellMethod,
+)
+
 from .sut_node import SutNode
 from .topology import Topology
 
@@ -96,6 +111,134 @@ def __hash__(self) -> int:
         """The subclasses must be hashable so that they can be stored in sets."""
 
 
+@dataclass
+class DecoratedNicCapability(Capability):
+    """A wrapper around :class:`~framework.remote_session.testpmd_shell.NicCapability`.
+
+    New instances should be created with the :meth:`create_unique` class method to ensure
+    there are no duplicate instances.
+
+    Attributes:
+        nic_capability: The NIC capability that defines each instance.
+        capability_fn: The capability retrieval function of `nic_capability`.
+        capability_decorator: The decorator function of `nic_capability`.
+            This function will wrap `capability_fn`.
+    """
+
+    nic_capability: NicCapability
+    capability_fn: TestPmdShellCapabilityMethod
+    capability_decorator: TestPmdShellDecorator | None
+    _unique_capabilities: ClassVar[dict[NicCapability, Self]] = {}
+
+    @classmethod
+    def get_unique(cls, nic_capability: NicCapability) -> "DecoratedNicCapability":
+        """Get the capability uniquely identified by `nic_capability`.
+
+        This is a factory method that implements a quasi-enum pattern.
+        The instances of this class are stored in an internal class variable,
+        `_unique_capabilities`.
+
+        If an instance identified by `nic_capability` doesn't exist,
+        it is created and added to `_unique_capabilities`.
+        If it exists, it is returned so that a new identical instance is not created.
+
+        Args:
+            nic_capability: The NIC capability.
+
+        Returns:
+            The capability uniquely identified by `nic_capability`.
+        """
+        decorator_fn = None
+        if isinstance(nic_capability.value, tuple):
+            capability_fn, decorator_fn = nic_capability.value
+        else:
+            capability_fn = nic_capability.value
+
+        if nic_capability not in cls._unique_capabilities:
+            cls._unique_capabilities[nic_capability] = cls(
+                nic_capability, capability_fn, decorator_fn
+            )
+        return cls._unique_capabilities[nic_capability]
+
+    @classmethod
+    def get_supported_capabilities(
+        cls, sut_node: SutNode, topology: "Topology"
+    ) -> set["DecoratedNicCapability"]:
+        """Overrides :meth:`~Capability.get_supported_capabilities`.
+
+        The capabilities are first sorted by decorators, then reduced into a single function which
+        is then passed to the decorator. This way we execute each decorator only once.
+        Each capability is first checked whether it's supported/unsupported
+        before executing its `capability_fn` so that each capability is retrieved only once.
+        """
+        supported_conditional_capabilities: set["DecoratedNicCapability"] = set()
+        logger = get_dts_logger(f"{sut_node.name}.{cls.__name__}")
+        if topology.type is topology.type.no_link:
+            logger.debug(
+                "No links available in the current topology, not getting NIC capabilities."
+            )
+            return supported_conditional_capabilities
+        logger.debug(
+            f"Checking which NIC capabilities from {cls.capabilities_to_check} are supported."
+        )
+        if cls.capabilities_to_check:
+            capabilities_to_check_map = cls._get_decorated_capabilities_map()
+            with TestPmdShell(
+                sut_node, privileged=True, disable_device_start=True
+            ) as testpmd_shell:
+                for conditional_capability_fn, capabilities in capabilities_to_check_map.items():
+                    supported_capabilities: set[NicCapability] = set()
+                    unsupported_capabilities: set[NicCapability] = set()
+                    capability_fn = cls._reduce_capabilities(
+                        capabilities, supported_capabilities, unsupported_capabilities
+                    )
+                    if conditional_capability_fn:
+                        capability_fn = conditional_capability_fn(capability_fn)
+                    capability_fn(testpmd_shell)
+                    for capability in capabilities:
+                        if capability.nic_capability in supported_capabilities:
+                            supported_conditional_capabilities.add(capability)
+
+        logger.debug(f"Found supported capabilities {supported_conditional_capabilities}.")
+        return supported_conditional_capabilities
+
+    @classmethod
+    def _get_decorated_capabilities_map(
+        cls,
+    ) -> dict[TestPmdShellDecorator | None, set["DecoratedNicCapability"]]:
+        capabilities_map: dict[TestPmdShellDecorator | None, set["DecoratedNicCapability"]] = {}
+        for capability in cls.capabilities_to_check:
+            if capability.capability_decorator not in capabilities_map:
+                capabilities_map[capability.capability_decorator] = set()
+            capabilities_map[capability.capability_decorator].add(capability)
+
+        return capabilities_map
+
+    @classmethod
+    def _reduce_capabilities(
+        cls,
+        capabilities: set["DecoratedNicCapability"],
+        supported_capabilities: MutableSet,
+        unsupported_capabilities: MutableSet,
+    ) -> TestPmdShellMethod:
+        def reduced_fn(testpmd_shell: TestPmdShell) -> None:
+            for capability in capabilities:
+                if capability not in supported_capabilities | unsupported_capabilities:
+                    capability.capability_fn(
+                        testpmd_shell, supported_capabilities, unsupported_capabilities
+                    )
+
+        return reduced_fn
+
+    def __hash__(self) -> int:
+        """Instances are identified by :attr:`nic_capability` and :attr:`capability_decorator`."""
+        return hash(self.nic_capability)
+
+    def __repr__(self) -> str:
+        """Easy to read string of :attr:`nic_capability` and :attr:`capability_decorator`."""
+        return f"{self.nic_capability}"
+
+
 class TestProtocol(Protocol):
     """Common test suite and test case attributes."""
 
@@ -116,6 +259,28 @@ def get_test_cases(cls, test_case_sublist: Sequence[str] | None = None) -> tuple
         raise NotImplementedError()
 
 
+def requires(
+    *nic_capabilities: NicCapability,
+) -> Callable[[type[TestProtocol]], type[TestProtocol]]:
+    """A decorator that adds the required capabilities to a test case or test suite.
+
+    Args:
+        nic_capabilities: The NIC capabilities that are required by the test case or test suite.
+
+    Returns:
+        The decorated test case or test suite.
+    """
+
+    def add_required_capability(test_case_or_suite: type[TestProtocol]) -> type[TestProtocol]:
+        for nic_capability in nic_capabilities:
+            decorated_nic_capability = DecoratedNicCapability.get_unique(nic_capability)
+            decorated_nic_capability.add_to_required(test_case_or_suite)
+
+        return test_case_or_suite
+
+    return add_required_capability
+
+
 def get_supported_capabilities(
     sut_node: SutNode,
     topology_config: Topology,
-- 
2.43.0