#2238 - Fixed the observations issue causing tests to fail

src/primaite/game/agent/observations.py

@@ -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": {}}
+            obs_dict = {"nic_status": 1 if nic_state["enabled"] else 2}
-            if CAPTURE_NMNE and nic_state.get("nmne"):
+            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)

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

@@ -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):

src/primaite/simulator/network/hardware/nodes/host/host_node.py

@@ -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