#2618 - Fixed the error that had NIC traffic describe state with enums as keys in the dict
This commit is contained in:
@@ -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.processes.process import Process
|
||||||
from primaite.simulator.system.services.service import Service
|
from primaite.simulator.system.services.service import Service
|
||||||
from primaite.simulator.system.software import IOSoftware
|
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
|
from primaite.utils.validators import IPV4Address
|
||||||
|
|
||||||
IOSoftwareClass = TypeVar("IOSoftwareClass", bound=IOSoftware)
|
IOSoftwareClass = TypeVar("IOSoftwareClass", bound=IOSoftware)
|
||||||
@@ -151,7 +152,7 @@ class NetworkInterface(SimComponent, ABC):
|
|||||||
)
|
)
|
||||||
if CAPTURE_NMNE:
|
if CAPTURE_NMNE:
|
||||||
state.update({"nmne": {k: v for k, v in self.nmne.items()}})
|
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
|
return state
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
|
|||||||
25
src/primaite/utils/converters.py
Normal file
25
src/primaite/utils/converters.py
Normal file
@@ -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
|
||||||
0
tests/unit_tests/_primaite/_utils/__init__.py
Normal file
0
tests/unit_tests/_primaite/_utils/__init__.py
Normal file
@@ -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
|
||||||
Reference in New Issue
Block a user