[PATCH 6/6] dts: restore TG setup and teardown

[Luca Vizzarro]

The setup and teardown was previously missed when reworking the
execution internals into states. Add back the 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 84d8fb26ac..774036571b 100644
--- a/dts/framework/test_run.py
+++ b/dts/framework/test_run.py
@@ -338,6 +338,7 @@ def next(self) -> State | None:
         test_run.ctx.sut_node.setup()
         test_run.ctx.tg_node.setup()
         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
@@ -420,6 +421,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.test_run.ctx.tg_node.teardown()
         self.test_run.ctx.sut_node.teardown()
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