[PATCH] dts: add TG setup and teardown

[Luca Vizzarro]

Add traffic generator setup and teardown respectively in
TestRunSetup and TestRunTeardown.

Signed-off-by: Luca Vizzarro <luca.vizzarro at arm.com>
Reviewed-by: Paul Szczepanek <paul.szczepanek at arm.com>
---
 dts/framework/test_run.py                                       | 2 ++
 .../testbed_model/traffic_generator/traffic_generator.py        | 2 +-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/dts/framework/test_run.py b/dts/framework/test_run.py
index 9f37c4bb6c..2808d013f6 100644
--- a/dts/framework/test_run.py
+++ b/dts/framework/test_run.py
@@ -334,6 +334,7 @@ def next(self) -> State | None:
         test_run.remaining_tests = deque(test_run.selected_tests)
 
         test_run.ctx.dpdk.setup(test_run.ctx.topology.sut_ports)
+        test_run.ctx.tg.setup(test_run.ctx.topology.tg_ports)
 
         self.result.ports = test_run.ctx.topology.sut_ports + test_run.ctx.topology.tg_ports
         self.result.sut_info = test_run.ctx.sut_node.node_info
@@ -416,6 +417,7 @@ def description(self) -> str:
 
     def next(self) -> State | None:
         """Next state."""
+        self.test_run.ctx.tg.teardown(self.test_run.ctx.topology.tg_ports)
         self.test_run.ctx.dpdk.teardown(self.test_run.ctx.topology.sut_ports)
         self.result.update_teardown(Result.PASS)
         return None
diff --git a/dts/framework/testbed_model/traffic_generator/traffic_generator.py b/dts/framework/testbed_model/traffic_generator/traffic_generator.py
index 4469273e36..804662e114 100644
--- a/dts/framework/testbed_model/traffic_generator/traffic_generator.py
+++ b/dts/framework/testbed_model/traffic_generator/traffic_generator.py
@@ -53,7 +53,7 @@ def __init__(self, tg_node: Node, config: TrafficGeneratorConfig, **kwargs):
     def setup(self, ports: Iterable[Port]):
         """Setup the traffic generator."""
 
-    def teardown(self):
+    def teardown(self, ports: Iterable[Port]):
         """Teardown the traffic generator."""
 
     def send_packet(self, packet: Packet, port: Port) -> None:
-- 
2.43.0