#2618 - Fixed the error that had NIC traffic describe state with enums as keys in the dict

src/primaite/simulator/network/hardware/base.py

@@ -37,6 +37,7 @@ from primaite.simulator.system.core.sys_log import SysLog
 from primaite.simulator.system.processes.process import Process
 from primaite.simulator.system.services.service import Service
 from primaite.simulator.system.software import IOSoftware
+from primaite.utils.converters import convert_dict_enum_keys_to_enum_values
 from primaite.utils.validators import IPV4Address
 
 IOSoftwareClass = TypeVar("IOSoftwareClass", bound=IOSoftware)
@@ -151,7 +152,7 @@ class NetworkInterface(SimComponent, ABC):
         )
         if CAPTURE_NMNE:
             state.update({"nmne": {k: v for k, v in self.nmne.items()}})
-        state.update({"traffic": self.traffic})
+        state.update({"traffic": convert_dict_enum_keys_to_enum_values(self.traffic)})
         return state
 
     @abstractmethod

src/primaite/utils/converters.py

@@ -0,0 +1,25 @@
+from enum import Enum
+from typing import Any, Dict
+
+
+def convert_dict_enum_keys_to_enum_values(d: Dict[Any, Any]) -> Dict[Any, Any]:
+    """
+    Convert dictionary keys from enums to their corresponding values.
+
+    :param d: dict
+        The dictionary with enum keys to be converted.
+    :return: dict
+        The dictionary with enum values as keys.
+    """
+    result = {}
+    for key, value in d.items():
+        if isinstance(key, Enum):
+            new_key = key.value
+        else:
+            new_key = key
+
+        if isinstance(value, dict):
+            result[new_key] = convert_dict_enum_keys_to_enum_values(value)
+        else:
+            result[new_key] = value
+    return result

tests/unit_tests/_primaite/_utils/test_dict_enum_keys_conversion.py

@@ -0,0 +1,83 @@
+from primaite.simulator.network.transmission.network_layer import IPProtocol
+from primaite.simulator.network.transmission.transport_layer import Port
+from primaite.utils.converters import convert_dict_enum_keys_to_enum_values
+
+
+def test_simple_conversion():
+    """
+    Test conversion of a simple dictionary with enum keys to enum values.
+
+    The original dictionary contains one level of nested dictionary with enums as keys.
+    The expected output should have string values of enums as keys.
+    """
+    original_dict = {IPProtocol.UDP: {Port.ARP: {"inbound": 0, "outbound": 1016.0}}}
+    expected_dict = {"udp": {219: {"inbound": 0, "outbound": 1016.0}}}
+    assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict
+
+
+def test_no_enums():
+    """
+    Test conversion of a dictionary with no enum keys.
+
+    The original dictionary contains only string keys.
+    The expected output should be identical to the original dictionary.
+    """
+    original_dict = {"protocol": {"port": {"inbound": 0, "outbound": 1016.0}}}
+    expected_dict = {"protocol": {"port": {"inbound": 0, "outbound": 1016.0}}}
+    assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict
+
+
+def test_mixed_keys():
+    """
+    Test conversion of a dictionary with a mix of enum and string keys.
+
+    The original dictionary contains both enums and strings as keys.
+    The expected output should have string values of enums and original string keys.
+    """
+    original_dict = {
+        IPProtocol.TCP: {"port": {"inbound": 0, "outbound": 1016.0}},
+        "protocol": {Port.HTTP: {"inbound": 10, "outbound": 2020.0}},
+    }
+    expected_dict = {
+        "tcp": {"port": {"inbound": 0, "outbound": 1016.0}},
+        "protocol": {80: {"inbound": 10, "outbound": 2020.0}},
+    }
+    assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict
+
+
+def test_empty_dict():
+    """
+    Test conversion of an empty dictionary.
+
+    The original dictionary is empty.
+    The expected output should also be an empty dictionary.
+    """
+    original_dict = {}
+    expected_dict = {}
+    assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict
+
+
+def test_nested_dicts():
+    """
+    Test conversion of a nested dictionary with multiple levels of nested dictionaries and enums as keys.
+
+    The original dictionary contains nested dictionaries with enums as keys at different levels.
+    The expected output should have string values of enums as keys at all levels.
+    """
+    original_dict = {
+        IPProtocol.UDP: {Port.ARP: {"inbound": 0, "outbound": 1016.0, "details": {IPProtocol.TCP: {"latency": "low"}}}}
+    }
+    expected_dict = {"udp": {219: {"inbound": 0, "outbound": 1016.0, "details": {"tcp": {"latency": "low"}}}}}
+    assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict
+
+
+def test_non_dict_values():
+    """
+    Test conversion of a dictionary where some values are not dictionaries.
+
+    The original dictionary contains lists and tuples as values.
+    The expected output should preserve these non-dictionary values while converting enum keys to string values.
+    """
+    original_dict = {IPProtocol.UDP: [Port.ARP, Port.HTTP], "protocols": (IPProtocol.TCP, IPProtocol.UDP)}
+    expected_dict = {"udp": [Port.ARP, Port.HTTP], "protocols": (IPProtocol.TCP, IPProtocol.UDP)}
+    assert convert_dict_enum_keys_to_enum_values(original_dict) == expected_dict