[PATCH v3 3/3] dts: resolve mypy type errors

[Paul Szczepanek <paul.szczepanek@arm.com>]

From: Luca Vizzarro <luca.vizzarro@arm.com>

The addition of scapy types yielded new errors in mypy, which could not
previously be checked.

Signed-off-by: Luca Vizzarro <luca.vizzarro@arm.com>
Reviewed-by: Paul Szczepanek <paul.szczepanek@arm.com>
---
 .../interactive_remote_session.py             |  4 ++--
 .../single_active_interactive_shell.py        |  2 +-
 dts/framework/remote_session/ssh_session.py   |  4 ++--
 dts/framework/test_suite.py                   | 20 +++++++++-------
 dts/framework/testbed_model/capability.py     |  4 ++--
 dts/framework/testbed_model/tg_node.py        |  2 +-
 .../capturing_traffic_generator.py            |  4 ++--
 .../testbed_model/traffic_generator/scapy.py  |  6 ++---
 .../traffic_generator/traffic_generator.py    |  2 +-
 dts/framework/utils.py                        |  6 ++---
 dts/tests/TestSuite_checksum_offload.py       | 24 +++++++++----------
 dts/tests/TestSuite_dynamic_queue_conf.py     | 15 ++++++++----
 dts/tests/TestSuite_mac_filter.py             |  6 ++---
 dts/tests/TestSuite_pmd_buffer_scatter.py     | 17 ++++++++-----
 dts/tests/TestSuite_vlan.py                   |  6 ++---
 15 files changed, 68 insertions(+), 54 deletions(-)

diff --git a/dts/framework/remote_session/interactive_remote_session.py b/dts/framework/remote_session/interactive_remote_session.py
index 4605ee14b4..509a284aaf 100644
--- a/dts/framework/remote_session/interactive_remote_session.py
+++ b/dts/framework/remote_session/interactive_remote_session.py
@@ -7,8 +7,8 @@
 import traceback
 from typing import Union
 
-from paramiko import AutoAddPolicy, SSHClient, Transport  # type: ignore[import-untyped]
-from paramiko.ssh_exception import (  # type: ignore[import-untyped]
+from paramiko import AutoAddPolicy, SSHClient, Transport
+from paramiko.ssh_exception import (
     AuthenticationException,
     BadHostKeyException,
     NoValidConnectionsError,
diff --git a/dts/framework/remote_session/single_active_interactive_shell.py b/dts/framework/remote_session/single_active_interactive_shell.py
index 513034d0e4..6de5fb33f5 100644
--- a/dts/framework/remote_session/single_active_interactive_shell.py
+++ b/dts/framework/remote_session/single_active_interactive_shell.py
@@ -24,7 +24,7 @@
 from pathlib import PurePath
 from typing import ClassVar
 
-from paramiko import Channel, channel  # type: ignore[import-untyped]
+from paramiko import Channel, channel
 from typing_extensions import Self
 
 from framework.exception import (
diff --git a/dts/framework/remote_session/ssh_session.py b/dts/framework/remote_session/ssh_session.py
index 329121913f..e6e4704bc2 100644
--- a/dts/framework/remote_session/ssh_session.py
+++ b/dts/framework/remote_session/ssh_session.py
@@ -8,12 +8,12 @@
 from pathlib import Path, PurePath
 
 from fabric import Connection  # type: ignore[import-untyped]
-from invoke.exceptions import (  # type: ignore[import-untyped]
+from invoke.exceptions import (
     CommandTimedOut,
     ThreadException,
     UnexpectedExit,
 )
-from paramiko.ssh_exception import (  # type: ignore[import-untyped]
+from paramiko.ssh_exception import (
     AuthenticationException,
     BadHostKeyException,
     NoValidConnectionsError,
diff --git a/dts/framework/test_suite.py b/dts/framework/test_suite.py
index 161bb10066..16012bfc79 100644
--- a/dts/framework/test_suite.py
+++ b/dts/framework/test_suite.py
@@ -26,9 +26,9 @@
 from types import ModuleType
 from typing import ClassVar, Protocol, TypeVar, Union, cast
 
-from scapy.layers.inet import IP  # type: ignore[import-untyped]
-from scapy.layers.l2 import Ether  # type: ignore[import-untyped]
-from scapy.packet import Packet, Padding, raw  # type: ignore[import-untyped]
+from scapy.layers.inet import IP
+from scapy.layers.l2 import Ether
+from scapy.packet import Packet, Padding, raw
 from typing_extensions import Self
 
 from framework.testbed_model.capability import TestProtocol
@@ -322,7 +322,7 @@ def _adjust_addresses(self, packets: list[Packet], expected: bool = False) -> li
         """
         ret_packets = []
         for original_packet in packets:
-            packet = original_packet.copy()
+            packet: Packet = original_packet.copy()
 
             # update l2 addresses
             # If `expected` is :data:`True`, the packet enters the TG from SUT, otherwise the
@@ -351,12 +351,13 @@ def _adjust_addresses(self, packets: list[Packet], expected: bool = False) -> li
                 # Update the last IP layer if there are multiple (the framework should be modifying
                 # the packet address instead of the tunnel address if there is one).
                 l3_to_use = packet.getlayer(IP, num_ip_layers)
+                assert l3_to_use is not None
                 if "src" not in l3_to_use.fields:
                     l3_to_use.src = self._tg_ip_address_egress.ip.exploded
 
                 if "dst" not in l3_to_use.fields:
                     l3_to_use.dst = self._tg_ip_address_ingress.ip.exploded
-            ret_packets.append(Ether(packet.build()))
+            ret_packets.append(packet)
 
         return ret_packets
 
@@ -405,7 +406,7 @@ def verify_packets(self, expected_packet: Packet, received_packets: list[Packet]
                 break
         else:
             self._logger.debug(
-                f"The expected packet {get_packet_summaries(expected_packet)} "
+                f"The expected packet {expected_packet.summary()} "
                 f"not found among received {get_packet_summaries(received_packets)}"
             )
             self._fail_test_case_verify("An expected packet not found among received packets.")
@@ -458,12 +459,13 @@ def _compare_packets(self, expected_packet: Packet, received_packet: Packet) ->
             self._logger.debug("Comparing payloads:")
             self._logger.debug(f"Received: {received_payload}")
             self._logger.debug(f"Expected: {expected_payload}")
-            if received_payload.__class__ == expected_payload.__class__:
+            if type(received_payload) is type(expected_payload):
                 self._logger.debug("The layers are the same.")
-                if received_payload.__class__ == Ether:
+                if type(received_payload) is Ether:
                     if not self._verify_l2_frame(received_payload, l3):
                         return False
-                elif received_payload.__class__ == IP:
+                elif type(received_payload) is IP:
+                    assert type(expected_payload) is IP
                     if not self._verify_l3_packet(received_payload, expected_payload):
                         return False
             else:
diff --git a/dts/framework/testbed_model/capability.py b/dts/framework/testbed_model/capability.py
index 63f99c4479..6a7a1f5b6c 100644
--- a/dts/framework/testbed_model/capability.py
+++ b/dts/framework/testbed_model/capability.py
@@ -169,7 +169,7 @@ class DecoratedNicCapability(Capability):
     _unique_capabilities: ClassVar[dict[NicCapability, Self]] = {}
 
     @classmethod
-    def get_unique(cls, nic_capability: NicCapability) -> "DecoratedNicCapability":
+    def get_unique(cls, nic_capability: NicCapability) -> Self:
         """Get the capability uniquely identified by `nic_capability`.
 
         This is a factory method that implements a quasi-enum pattern.
@@ -314,7 +314,7 @@ def _preprocess_required(self, test_case_or_suite: type["TestProtocol"]) -> None
         test_case_or_suite.topology_type = self
 
     @classmethod
-    def get_unique(cls, topology_type: TopologyType) -> "TopologyCapability":
+    def get_unique(cls, topology_type: TopologyType) -> Self:
         """Get the capability uniquely identified by `topology_type`.
 
         This is a factory method that implements a quasi-enum pattern.
diff --git a/dts/framework/testbed_model/tg_node.py b/dts/framework/testbed_model/tg_node.py
index 4179365abb..3071bbd645 100644
--- a/dts/framework/testbed_model/tg_node.py
+++ b/dts/framework/testbed_model/tg_node.py
@@ -9,7 +9,7 @@
 A TG node is where the TG runs.
 """
 
-from scapy.packet import Packet  # type: ignore[import-untyped]
+from scapy.packet import Packet
 
 from framework.config import TGNodeConfiguration
 from framework.testbed_model.traffic_generator.capturing_traffic_generator import (
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 66a77da9c4..e31ba2a9b7 100644
--- a/dts/framework/testbed_model/traffic_generator/capturing_traffic_generator.py
+++ b/dts/framework/testbed_model/traffic_generator/capturing_traffic_generator.py
@@ -13,8 +13,8 @@
 from abc import abstractmethod
 from dataclasses import dataclass
 
-import scapy.utils  # type: ignore[import-untyped]
-from scapy.packet import Packet  # type: ignore[import-untyped]
+import scapy.utils
+from scapy.packet import Packet
 
 from framework.settings import SETTINGS
 from framework.testbed_model.port import Port
diff --git a/dts/framework/testbed_model/traffic_generator/scapy.py b/dts/framework/testbed_model/traffic_generator/scapy.py
index 0084959d57..d24efc44e6 100644
--- a/dts/framework/testbed_model/traffic_generator/scapy.py
+++ b/dts/framework/testbed_model/traffic_generator/scapy.py
@@ -16,9 +16,9 @@
 import time
 from typing import ClassVar
 
-from scapy.compat import base64_bytes  # type: ignore[import-untyped]
-from scapy.layers.l2 import Ether  # type: ignore[import-untyped]
-from scapy.packet import Packet  # type: ignore[import-untyped]
+from scapy.compat import base64_bytes
+from scapy.layers.l2 import Ether
+from scapy.packet import Packet
 
 from framework.config import OS, ScapyTrafficGeneratorConfig
 from framework.remote_session.python_shell import PythonShell
diff --git a/dts/framework/testbed_model/traffic_generator/traffic_generator.py b/dts/framework/testbed_model/traffic_generator/traffic_generator.py
index 5aeb9147c0..a07538cc98 100644
--- a/dts/framework/testbed_model/traffic_generator/traffic_generator.py
+++ b/dts/framework/testbed_model/traffic_generator/traffic_generator.py
@@ -10,7 +10,7 @@
 
 from abc import ABC, abstractmethod
 
-from scapy.packet import Packet  # type: ignore[import-untyped]
+from scapy.packet import Packet
 
 from framework.config import TrafficGeneratorConfig
 from framework.logger import DTSLogger, get_dts_logger
diff --git a/dts/framework/utils.py b/dts/framework/utils.py
index 6839bcf243..66f37a8813 100644
--- a/dts/framework/utils.py
+++ b/dts/framework/utils.py
@@ -23,8 +23,8 @@
 from pathlib import Path
 from typing import Any, Callable
 
-from scapy.layers.inet import IP, TCP, UDP, Ether  # type: ignore[import-untyped]
-from scapy.packet import Packet  # type: ignore[import-untyped]
+from scapy.layers.inet import IP, TCP, UDP, Ether
+from scapy.packet import Packet
 
 from .exception import InternalError
 
@@ -270,7 +270,7 @@ def generate_random_packets(
     if payload_size < 0:
         raise InternalError(f"An invalid payload_size of {payload_size} was given.")
 
-    l4_factories = []
+    l4_factories: list[type[Packet]] = []
     if protocols & PacketProtocols.TCP:
         l4_factories.append(TCP)
     if protocols & PacketProtocols.UDP:
diff --git a/dts/tests/TestSuite_checksum_offload.py b/dts/tests/TestSuite_checksum_offload.py
index 1a933563b6..c1680bd388 100644
--- a/dts/tests/TestSuite_checksum_offload.py
+++ b/dts/tests/TestSuite_checksum_offload.py
@@ -13,11 +13,11 @@
 
 from typing import List
 
-from scapy.layers.inet import IP, TCP, UDP  # type: ignore[import-untyped]
-from scapy.layers.inet6 import IPv6  # type: ignore[import-untyped]
-from scapy.layers.l2 import Dot1Q, Ether  # type: ignore[import-untyped]
-from scapy.layers.sctp import SCTP  # type: ignore[import-untyped]
-from scapy.packet import Packet, Raw  # type: ignore[import-untyped]
+from scapy.layers.inet import IP, TCP, UDP
+from scapy.layers.inet6 import IPv6
+from scapy.layers.l2 import Dot1Q, Ether
+from scapy.layers.sctp import SCTP
+from scapy.packet import Packet, Raw
 
 from framework.params.testpmd import SimpleForwardingModes
 from framework.remote_session.testpmd_shell import (
@@ -49,7 +49,7 @@ class TestChecksumOffload(TestSuite):
     """
 
     def send_packets_and_verify(
-        self, packet_list: List[Packet], load: str, should_receive: bool
+        self, packet_list: List[Packet], load: bytes, should_receive: bool
     ) -> None:
         """Iterates through a list of packets and verifies they are received.
 
@@ -62,7 +62,7 @@ def send_packets_and_verify(
         for i in range(0, len(packet_list)):
             received_packets = self.send_packet_and_capture(packet=packet_list[i])
             received = any(
-                packet.haslayer(Raw) and load in str(packet.load) for packet in received_packets
+                packet.haslayer(Raw) and load in packet.load for packet in received_packets
             )
             self.verify(
                 received == should_receive,
@@ -86,8 +86,8 @@ def send_packet_and_verify_checksum(
         testpmd.start()
         self.send_packet_and_capture(packet=packet)
         verbose_output = testpmd.extract_verbose_output(testpmd.stop())
-        for packet in verbose_output:
-            if packet.dst_mac == id:
+        for testpmd_packet in verbose_output:
+            if testpmd_packet.dst_mac == id:
                 isIP = PacketOffloadFlag.RTE_MBUF_F_RX_IP_CKSUM_GOOD in packet.ol_flags
                 isL4 = PacketOffloadFlag.RTE_MBUF_F_RX_L4_CKSUM_GOOD in packet.ol_flags
         self.verify(isL4 == goodL4, "Layer 4 checksum flag did not match expected checksum flag.")
@@ -110,7 +110,7 @@ def setup_hw_offload(self, testpmd: TestPmdShell) -> None:
     def test_insert_checksums(self) -> None:
         """Enable checksum offload insertion and verify packet reception."""
         mac_id = "00:00:00:00:00:01"
-        payload = "xxxxx"
+        payload = b"xxxxx"
         packet_list = [
             Ether(dst=mac_id) / IP() / UDP() / Raw(payload),
             Ether(dst=mac_id) / IP() / TCP() / Raw(payload),
@@ -132,7 +132,7 @@ def test_insert_checksums(self) -> None:
     def test_no_insert_checksums(self) -> None:
         """Disable checksum offload insertion and verify packet reception."""
         mac_id = "00:00:00:00:00:01"
-        payload = "xxxxx"
+        payload = b"xxxxx"
         packet_list = [
             Ether(dst=mac_id) / IP() / UDP() / Raw(payload),
             Ether(dst=mac_id) / IP() / TCP() / Raw(payload),
@@ -231,7 +231,7 @@ def test_validate_rx_checksum(self) -> None:
     def test_vlan_checksum(self) -> None:
         """Test VLAN Rx checksum hardware offload and verify packet reception."""
         mac_id = "00:00:00:00:00:01"
-        payload = "xxxxx"
+        payload = b"xxxxx"
         packet_list = [
             Ether(dst=mac_id) / Dot1Q(vlan=1) / IP(chksum=0x0) / UDP(chksum=0xF) / Raw(payload),
             Ether(dst=mac_id) / Dot1Q(vlan=1) / IP(chksum=0x0) / TCP(chksum=0xF) / Raw(payload),
diff --git a/dts/tests/TestSuite_dynamic_queue_conf.py b/dts/tests/TestSuite_dynamic_queue_conf.py
index ac7b81b602..e55716f545 100644
--- a/dts/tests/TestSuite_dynamic_queue_conf.py
+++ b/dts/tests/TestSuite_dynamic_queue_conf.py
@@ -27,9 +27,9 @@
 import random
 from typing import Callable, ClassVar, MutableSet
 
-from scapy.layers.inet import IP  # type: ignore[import-untyped]
-from scapy.layers.l2 import Ether  # type: ignore[import-untyped]
-from scapy.packet import Raw  # type: ignore[import-untyped]
+from scapy.layers.inet import IP
+from scapy.layers.l2 import Ether
+from scapy.packet import Raw
 
 from framework.exception import InteractiveCommandExecutionError
 from framework.params.testpmd import PortTopology, SimpleForwardingModes
@@ -93,7 +93,14 @@ def wrap(self: "TestDynamicQueueConf", is_rx_testing: bool) -> None:
                 testpmd.stop_port_queue(port_id, q, is_rx_testing)
             testpmd.set_forward_mode(SimpleForwardingModes.mac)
 
-            test_meth(self, port_id, queues_to_config, unchanged_queues, testpmd, is_rx_testing)
+            test_meth(
+                self,
+                port_id,
+                queues_to_config,
+                unchanged_queues,
+                testpmd,
+                is_rx_testing,
+            )
 
             for queue_id in queues_to_config:
                 testpmd.start_port_queue(port_id, queue_id, is_rx_testing)
diff --git a/dts/tests/TestSuite_mac_filter.py b/dts/tests/TestSuite_mac_filter.py
index cc9e9fd6b4..11e4b595c7 100644
--- a/dts/tests/TestSuite_mac_filter.py
+++ b/dts/tests/TestSuite_mac_filter.py
@@ -15,9 +15,9 @@
 allow list is behaving as expected.
 """
 
-from scapy.layers.inet import IP  # type: ignore[import-untyped]
-from scapy.layers.l2 import Ether  # type: ignore[import-untyped]
-from scapy.packet import Raw  # type: ignore[import-untyped]
+from scapy.layers.inet import IP
+from scapy.layers.l2 import Ether
+from scapy.packet import Raw
 
 from framework.exception import InteractiveCommandExecutionError
 from framework.remote_session.testpmd_shell import NicCapability, TestPmdShell
diff --git a/dts/tests/TestSuite_pmd_buffer_scatter.py b/dts/tests/TestSuite_pmd_buffer_scatter.py
index 6367612498..b2f42425d4 100644
--- a/dts/tests/TestSuite_pmd_buffer_scatter.py
+++ b/dts/tests/TestSuite_pmd_buffer_scatter.py
@@ -17,10 +17,10 @@
 
 import struct
 
-from scapy.layers.inet import IP  # type: ignore[import-untyped]
-from scapy.layers.l2 import Ether  # type: ignore[import-untyped]
-from scapy.packet import Raw  # type: ignore[import-untyped]
-from scapy.utils import hexstr  # type: ignore[import-untyped]
+from scapy.layers.inet import IP
+from scapy.layers.l2 import Ether
+from scapy.packet import Raw
+from scapy.utils import hexstr
 
 from framework.params.testpmd import SimpleForwardingModes
 from framework.remote_session.testpmd_shell import TestPmdShell
@@ -76,7 +76,8 @@ def scatter_pktgen_send_packet(self, pktsize: int) -> str:
             The payload of the received packet as a string.
         """
         packet = Ether() / IP() / Raw()
-        packet.getlayer(2).load = ""
+        if layer2 := packet.getlayer(2):
+            layer2.load = ""
         payload_len = pktsize - len(packet) - 4
         payload = ["58"] * payload_len
         # pack the payload
@@ -84,7 +85,11 @@ def scatter_pktgen_send_packet(self, pktsize: int) -> str:
             packet.load += struct.pack("=B", int("%s%s" % (X_in_hex[0], X_in_hex[1]), 16))
         received_packets = self.send_packet_and_capture(packet)
         self.verify(len(received_packets) > 0, "Did not receive any packets.")
-        load = hexstr(received_packets[0].getlayer(2), onlyhex=1)
+
+        layer2 = received_packets[0].getlayer(2)
+        self.verify(layer2 is not None, "The received packet is invalid.")
+        assert layer2 is not None
+        load = hexstr(layer2, onlyhex=1)
 
         return load
 
diff --git a/dts/tests/TestSuite_vlan.py b/dts/tests/TestSuite_vlan.py
index 35221fe362..c67520baef 100644
--- a/dts/tests/TestSuite_vlan.py
+++ b/dts/tests/TestSuite_vlan.py
@@ -12,8 +12,8 @@
 
 """
 
-from scapy.layers.l2 import Dot1Q, Ether  # type: ignore[import-untyped]
-from scapy.packet import Raw  # type: ignore[import-untyped]
+from scapy.layers.l2 import Dot1Q, Ether
+from scapy.packet import Raw
 
 from framework.remote_session.testpmd_shell import SimpleForwardingModes, TestPmdShell
 from framework.test_suite import TestSuite, func_test
@@ -96,7 +96,7 @@ def send_packet_and_verify_insertion(self, expected_id: int) -> None:
         )
         if test_packet is not None:
             self.verify(
-                test_packet.haslayer(Dot1Q),
+                test_packet.haslayer(Dot1Q) == 1,
                 "The received packet did not have a VLAN tag",
             )
             self.verify(
-- 
2.39.2