diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py index 0cbfcbb2..e9560b91 100644 --- a/src/primaite/simulator/network/hardware/base.py +++ b/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 diff --git a/src/primaite/utils/converters.py b/src/primaite/utils/converters.py new file mode 100644 index 00000000..506ef31f --- /dev/null +++ b/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 diff --git a/tests/unit_tests/_primaite/_utils/__init__.py b/tests/unit_tests/_primaite/_utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/_primaite/_utils/test_dict_enum_keys_conversion.py b/tests/unit_tests/_primaite/_utils/test_dict_enum_keys_conversion.py new file mode 100644 index 00000000..3dcc8be5 --- /dev/null +++ b/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