[PATCH v3 2/5] dts: rework traffic generator inheritance structure

[Patrick Robb]

From: Nicholas Pratte <npratte at iol.unh.edu>

Rework TG class hierarchy to include performance traffic generators.
As such, methods specific to capturing traffic have been moved to the
CapturingTrafficGenerator subclass.

Bugzilla ID: 1697
Signed-off-by: Nicholas Pratte <npratte at iol.unh.edu>
Signed-off-by: Patrick Robb <probb at iol.unh.edu>
Reviewed-by: Dean Marx <dmarx at iol.unh.edu>
---
 .../capturing_traffic_generator.py            | 34 ++++++++++
 .../performance_traffic_generator.py          | 63 +++++++++++++++++++
 .../testbed_model/traffic_generator/scapy.py  | 10 ++-
 .../traffic_generator/traffic_generator.py    | 44 -------------
 4 files changed, 104 insertions(+), 47 deletions(-)
 create mode 100644 dts/framework/testbed_model/traffic_generator/performance_traffic_generator.py

diff --git a/dts/framework/testbed_model/traffic_generator/capturing_traffic_generator.py b/dts/framework/testbed_model/traffic_generator/capturing_traffic_generator.py
index a85858ba07..110244a7b0 100644
--- a/dts/framework/testbed_model/traffic_generator/capturing_traffic_generator.py
+++ b/dts/framework/testbed_model/traffic_generator/capturing_traffic_generator.py
@@ -65,6 +65,40 @@ def is_capturing(self) -> bool:
         """This traffic generator can capture traffic."""
         return True
 
+    def send_packet(self, packet: Packet, port: Port) -> None:
+        """Send `packet` and block until it is fully sent.
+
+        Send `packet` on `port`, then wait until `packet` is fully sent.
+
+        Args:
+            packet: The packet to send.
+            port: The egress port on the TG node.
+        """
+        self.send_packets([packet], port)
+
+    def send_packets(self, packets: list[Packet], port: Port) -> None:
+        """Send `packets` and block until they are fully sent.
+
+        Send `packets` on `port`, then wait until `packets` are fully sent.
+
+        Args:
+            packets: The packets to send.
+            port: The egress port on the TG node.
+        """
+        self._logger.info(f"Sending packet{'s' if len(packets) > 1 else ''}.")
+        self._logger.debug(get_packet_summaries(packets))
+        self._send_packets(packets, port)
+
+    @abstractmethod
+    def _send_packets(self, packets: list[Packet], port: Port) -> None:
+        """The implementation of :method:`send_packets`.
+
+        The subclasses must implement this method which sends `packets` on `port`.
+        The method should block until all `packets` are fully sent.
+
+        What fully sent means is defined by the traffic generator.
+        """
+
     def send_packets_and_capture(
         self,
         packets: list[Packet],
diff --git a/dts/framework/testbed_model/traffic_generator/performance_traffic_generator.py b/dts/framework/testbed_model/traffic_generator/performance_traffic_generator.py
new file mode 100644
index 0000000000..59dac6a075
--- /dev/null
+++ b/dts/framework/testbed_model/traffic_generator/performance_traffic_generator.py
@@ -0,0 +1,63 @@
+"""Traffic generators for performance tests which can generate a high number of packets."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+from scapy.packet import Packet
+
+from framework.testbed_model.topology import Topology
+
+from .traffic_generator import TrafficGenerator
+
+
+ at dataclass(slots=True)
+class PerformanceTrafficStats(ABC):
+    """Data structure to store performance statistics for a given test run.
+
+    Attributes:
+        frame_size: The total length of the frame
+        tx_pps: Recorded tx packets per second
+        tx_cps: Recorded tx connections per second
+        tx_bps: Recorded tx bytes per second
+        rx_pps: Recorded rx packets per second
+        rx_bps: Recorded rx bytes per second
+    """
+
+    frame_size: int
+
+    tx_pps: float
+    tx_cps: float
+    tx_bps: float
+
+    rx_pps: float
+    rx_bps: float
+
+
+class PerformanceTrafficGenerator(TrafficGenerator):
+    """An abstract base class for all performance-oriented traffic generators.
+
+    Provides an intermediary interface for performance-based traffic generator.
+    """
+
+    @abstractmethod
+    def calculate_traffic_and_stats(
+        self,
+        packet: Packet,
+        send_mpps: int,
+        duration: float,
+    ) -> PerformanceTrafficStats:
+        """Send packet traffic and acquire associated statistics.
+
+        Args:
+        packet: The packet to send.
+        send_mpps: The millions packets per second send rate.
+        duration: Performance test duration (in seconds).
+
+        Returns:
+            Performance statistics of the generated test.
+        """
+
+    def setup(self, topology: Topology) -> None:
+        """Overrides :meth:`.traffic_generator.TrafficGenerator.setup`."""
+        for port in self._tg_node.ports:
+            self._tg_node.main_session.configure_port_mtu(2000, port)
diff --git a/dts/framework/testbed_model/traffic_generator/scapy.py b/dts/framework/testbed_model/traffic_generator/scapy.py
index e21ba4ed96..602b93d473 100644
--- a/dts/framework/testbed_model/traffic_generator/scapy.py
+++ b/dts/framework/testbed_model/traffic_generator/scapy.py
@@ -26,7 +26,7 @@
 from scapy.packet import Packet
 
 from framework.config.node import OS
-from framework.config.test_run import ScapyTrafficGeneratorConfig
+from framework.config.test_run import TrafficGeneratorConfig
 from framework.exception import InteractiveSSHSessionDeadError, InternalError
 from framework.remote_session.python_shell import PythonShell
 from framework.testbed_model.node import Node
@@ -285,7 +285,7 @@ class also extends :class:`.capturing_traffic_generator.CapturingTrafficGenerato
     first.
     """
 
-    _config: ScapyTrafficGeneratorConfig
+    _config: TrafficGeneratorConfig
     _shell: PythonShell
     _sniffer: ScapyAsyncSniffer
 
@@ -296,7 +296,7 @@ class also extends :class:`.capturing_traffic_generator.CapturingTrafficGenerato
     #: Padding to add to the start of a line for python syntax compliance.
     _python_indentation: ClassVar[str] = " " * 4
 
-    def __init__(self, tg_node: Node, config: ScapyTrafficGeneratorConfig, **kwargs):
+    def __init__(self, tg_node: Node, config: TrafficGeneratorConfig, **kwargs):
         """Extend the constructor with Scapy TG specifics.
 
         Initializes both the traffic generator and the interactive shell used to handle Scapy
@@ -332,6 +332,10 @@ def setup(self, topology: Topology):
         self._shell.send_command("from scapy.all import *")
         self._shell.send_command("from scapy.contrib.lldp import *")
 
+    def teardown(self):
+        """Overrides :meth:`.traffic_generator.TrafficGenerator.teardown`."""
+        pass
+
     def close(self):
         """Overrides :meth:`.traffic_generator.TrafficGenerator.close`.
 
diff --git a/dts/framework/testbed_model/traffic_generator/traffic_generator.py b/dts/framework/testbed_model/traffic_generator/traffic_generator.py
index 8f53b07daf..ea3075989d 100644
--- a/dts/framework/testbed_model/traffic_generator/traffic_generator.py
+++ b/dts/framework/testbed_model/traffic_generator/traffic_generator.py
@@ -10,14 +10,10 @@
 
 from abc import ABC, abstractmethod
 
-from scapy.packet import Packet
-
 from framework.config.test_run import TrafficGeneratorConfig
 from framework.logger import DTSLogger, get_dts_logger
 from framework.testbed_model.node import Node
-from framework.testbed_model.port import Port
 from framework.testbed_model.topology import Topology
-from framework.utils import get_packet_summaries
 
 
 class TrafficGenerator(ABC):
@@ -54,46 +50,6 @@ def setup(self, topology: Topology):
 
     def teardown(self):
         """Teardown the traffic generator."""
-        self.close()
-
-    def send_packet(self, packet: Packet, port: Port) -> None:
-        """Send `packet` and block until it is fully sent.
-
-        Send `packet` on `port`, then wait until `packet` is fully sent.
-
-        Args:
-            packet: The packet to send.
-            port: The egress port on the TG node.
-        """
-        self.send_packets([packet], port)
-
-    def send_packets(self, packets: list[Packet], port: Port) -> None:
-        """Send `packets` and block until they are fully sent.
-
-        Send `packets` on `port`, then wait until `packets` are fully sent.
-
-        Args:
-            packets: The packets to send.
-            port: The egress port on the TG node.
-        """
-        self._logger.info(f"Sending packet{'s' if len(packets) > 1 else ''}.")
-        self._logger.debug(get_packet_summaries(packets))
-        self._send_packets(packets, port)
-
-    @abstractmethod
-    def _send_packets(self, packets: list[Packet], port: Port) -> None:
-        """The implementation of :method:`send_packets`.
-
-        The subclasses must implement this method which sends `packets` on `port`.
-        The method should block until all `packets` are fully sent.
-
-        What fully sent means is defined by the traffic generator.
-        """
-
-    @property
-    def is_capturing(self) -> bool:
-        """This traffic generator can't capture traffic."""
-        return False
 
     @abstractmethod
     def close(self) -> None:
-- 
2.49.0