#2238 - Fixed the observations issue causing tests to fail

This commit is contained in:
Chris McCarthy
2024-02-28 12:03:58 +00:00
parent 52677538a8
commit d55b6a5b48
3 changed files with 8 additions and 3 deletions

View File

@@ -351,6 +351,8 @@ class NicObservation(AbstractObservation):
def default_observation(self) -> Dict:
"""The default NIC observation dict."""
data = {"nic_status": 0}
if CAPTURE_NMNE:
data.update({"nmne": {"inbound": 0, "outbound": 0}})
return data
@@ -404,8 +406,9 @@ class NicObservation(AbstractObservation):
if nic_state is NOT_PRESENT_IN_STATE:
return self.default_observation
else:
obs_dict = {"nic_status": 1 if nic_state["enabled"] else 2, "nmne": {}}
if CAPTURE_NMNE and nic_state.get("nmne"):
obs_dict = {"nic_status": 1 if nic_state["enabled"] else 2}
if CAPTURE_NMNE:
obs_dict.update({"nmne": {}})
direction_dict = nic_state["nmne"].get("direction", {})
inbound_keywords = direction_dict.get("inbound", {}).get("keywords", {})
inbound_count = inbound_keywords.get("*", 0)

View File

@@ -123,6 +123,8 @@ class NetworkInterface(SimComponent, ABC):
"enabled": self.enabled,
}
)
if CAPTURE_NMNE:
state.update({"nmne": self.nmne})
return state
def reset_component_for_episode(self, episode: int):

View File

@@ -205,7 +205,7 @@ class NIC(IPWiredNetworkInterface):
state = super().describe_state()
# Update the state with NIC-specific information
state.update({"wake_on_lan": self.wake_on_lan, "nmne": self.nmne})
state.update({"wake_on_lan": self.wake_on_lan})
return state