#2350: splitting observations into separate files

src/primaite/game/agent/interface.py

@@ -6,7 +6,7 @@ from gymnasium.core import ActType, ObsType
 from pydantic import BaseModel, model_validator
 
 from primaite.game.agent.actions import ActionManager
-from primaite.game.agent.observations import ObservationManager
+from primaite.game.agent.observations.observation_manager import ObservationManager
 from primaite.game.agent.rewards import RewardFunction
 
 if TYPE_CHECKING:
@@ -146,23 +146,10 @@ class AbstractAgent(ABC):
 class AbstractScriptedAgent(AbstractAgent):
     """Base class for actors which generate their own behaviour."""
 
-    pass
+    @abstractmethod
-
-
-class RandomAgent(AbstractScriptedAgent):
-    """Agent that ignores its observation and acts completely at random."""
-
     def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
-        """Sample the action space randomly.
+        """Return an action to be taken in the environment."""
-
+        return super().get_action(obs=obs, timestep=timestep)
-        :param obs: Current observation for this agent, not used in RandomAgent
-        :type obs: ObsType
-        :param timestep: The current simulation timestep, not used in RandomAgent
-        :type timestep: int
-        :return: Action formatted in CAOS format
-        :rtype: Tuple[str, Dict]
-        """
-        return self.action_manager.get_action(self.action_manager.space.sample())
 
 
 class ProxyAgent(AbstractAgent):

src/primaite/game/agent/observations/agent_observations.py

@@ -0,0 +1,188 @@
+from typing import Dict, List, Optional, Tuple, TYPE_CHECKING
+
+from gymnasium import spaces
+
+from primaite.game.agent.observations.observations import (
+    AbstractObservation,
+    AclObservation,
+    ICSObservation,
+    LinkObservation,
+    NodeObservation,
+    NullObservation,
+)
+
+if TYPE_CHECKING:
+    from primaite.game.game import PrimaiteGame
+
+
+class UC2BlueObservation(AbstractObservation):
+    """Container for all observations used by the blue agent in UC2.
+
+    TODO: there's no real need for a UC2 blue container class, we should be able to simply use the observation handler
+        for the purpose of compiling several observation components.
+    """
+
+    def __init__(
+        self,
+        nodes: List[NodeObservation],
+        links: List[LinkObservation],
+        acl: AclObservation,
+        ics: ICSObservation,
+        where: Optional[List[str]] = None,
+    ) -> None:
+        """Initialise UC2 blue observation.
+
+        :param nodes: List of node observations
+        :type nodes: List[NodeObservation]
+        :param links: List of link observations
+        :type links: List[LinkObservation]
+        :param acl: The Access Control List observation
+        :type acl: AclObservation
+        :param ics: The ICS observation
+        :type ics: ICSObservation
+        :param where: Where in the simulation state dict to find information. Not used in this particular observation
+            because it only compiles other observations and doesn't contribute any new information, defaults to None
+        :type where: Optional[List[str]], optional
+        """
+        super().__init__()
+        self.where: Optional[Tuple[str]] = where
+
+        self.nodes: List[NodeObservation] = nodes
+        self.links: List[LinkObservation] = links
+        self.acl: AclObservation = acl
+        self.ics: ICSObservation = ics
+
+        self.default_observation: Dict = {
+            "NODES": {i + 1: n.default_observation for i, n in enumerate(self.nodes)},
+            "LINKS": {i + 1: l.default_observation for i, l in enumerate(self.links)},
+            "ACL": self.acl.default_observation,
+            "ICS": self.ics.default_observation,
+        }
+
+    def observe(self, state: Dict) -> Dict:
+        """Generate observation based on the current state of the simulation.
+
+        :param state: Simulation state dictionary
+        :type state: Dict
+        :return: Observation
+        :rtype: Dict
+        """
+        if self.where is None:
+            return self.default_observation
+
+        obs = {}
+        obs["NODES"] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)}
+        obs["LINKS"] = {i + 1: link.observe(state) for i, link in enumerate(self.links)}
+        obs["ACL"] = self.acl.observe(state)
+        obs["ICS"] = self.ics.observe(state)
+
+        return obs
+
+    @property
+    def space(self) -> spaces.Space:
+        """
+        Gymnasium space object describing the observation space shape.
+
+        :return: Space
+        :rtype: spaces.Space
+        """
+        return spaces.Dict(
+            {
+                "NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}),
+                "LINKS": spaces.Dict({i + 1: link.space for i, link in enumerate(self.links)}),
+                "ACL": self.acl.space,
+                "ICS": self.ics.space,
+            }
+        )
+
+    @classmethod
+    def from_config(cls, config: Dict, game: "PrimaiteGame") -> "UC2BlueObservation":
+        """Create UC2 blue observation from a config.
+
+        :param config: Dictionary containing the configuration for this UC2 blue observation. This includes the nodes,
+            links, ACL and ICS observations.
+        :type config: Dict
+        :param game: Reference to the PrimaiteGame object that spawned this observation.
+        :type game: PrimaiteGame
+        :return: Constructed UC2 blue observation
+        :rtype: UC2BlueObservation
+        """
+        node_configs = config["nodes"]
+
+        num_services_per_node = config["num_services_per_node"]
+        num_folders_per_node = config["num_folders_per_node"]
+        num_files_per_folder = config["num_files_per_folder"]
+        num_nics_per_node = config["num_nics_per_node"]
+        nodes = [
+            NodeObservation.from_config(
+                config=n,
+                game=game,
+                num_services_per_node=num_services_per_node,
+                num_folders_per_node=num_folders_per_node,
+                num_files_per_folder=num_files_per_folder,
+                num_nics_per_node=num_nics_per_node,
+            )
+            for n in node_configs
+        ]
+
+        link_configs = config["links"]
+        links = [LinkObservation.from_config(config=link, game=game) for link in link_configs]
+
+        acl_config = config["acl"]
+        acl = AclObservation.from_config(config=acl_config, game=game)
+
+        ics_config = config["ics"]
+        ics = ICSObservation.from_config(config=ics_config, game=game)
+        new = cls(nodes=nodes, links=links, acl=acl, ics=ics, where=["network"])
+        return new
+
+
+class UC2RedObservation(AbstractObservation):
+    """Container for all observations used by the red agent in UC2."""
+
+    def __init__(self, nodes: List[NodeObservation], where: Optional[List[str]] = None) -> None:
+        super().__init__()
+        self.where: Optional[List[str]] = where
+        self.nodes: List[NodeObservation] = nodes
+
+        self.default_observation: Dict = {
+            "NODES": {i + 1: n.default_observation for i, n in enumerate(self.nodes)},
+        }
+
+    def observe(self, state: Dict) -> Dict:
+        """Generate observation based on the current state of the simulation."""
+        if self.where is None:
+            return self.default_observation
+
+        obs = {}
+        obs["NODES"] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)}
+        return obs
+
+    @property
+    def space(self) -> spaces.Space:
+        """Gymnasium space object describing the observation space shape."""
+        return spaces.Dict(
+            {
+                "NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}),
+            }
+        )
+
+    @classmethod
+    def from_config(cls, config: Dict, game: "PrimaiteGame") -> "UC2RedObservation":
+        """
+        Create UC2 red observation from a config.
+
+        :param config: Dictionary containing the configuration for this UC2 red observation.
+        :type config: Dict
+        :param game: Reference to the PrimaiteGame object that spawned this observation.
+        :type game: PrimaiteGame
+        """
+        node_configs = config["nodes"]
+        nodes = [NodeObservation.from_config(config=cfg, game=game) for cfg in node_configs]
+        return cls(nodes=nodes, where=["network"])
+
+
+class UC2GreenObservation(NullObservation):
+    """Green agent observation. As the green agent's actions don't depend on the observation, this is empty."""
+
+    pass

src/primaite/game/agent/observations/observation_manager.py

@@ -0,0 +1,73 @@
+from typing import Dict, TYPE_CHECKING
+
+from gymnasium.core import ObsType
+
+from primaite.game.agent.observations.agent_observations import (
+    UC2BlueObservation,
+    UC2GreenObservation,
+    UC2RedObservation,
+)
+from primaite.game.agent.observations.observations import AbstractObservation
+
+if TYPE_CHECKING:
+    from primaite.game.game import PrimaiteGame
+
+
+class ObservationManager:
+    """
+    Manage the observations of an Agent.
+
+    The observation space has the purpose of:
+      1. Reading the outputted state from the PrimAITE Simulation.
+      2. Selecting parts of the simulation state that are requested by the simulation config
+      3. Formatting this information so an agent can use it to make decisions.
+    """
+
+    # TODO: Dear code reader: This class currently doesn't do much except hold an observation object. It will be changed
+    # to have more of it's own behaviour, and it will replace UC2BlueObservation and UC2RedObservation during the next
+    # refactor.
+
+    def __init__(self, observation: AbstractObservation) -> None:
+        """Initialise observation space.
+
+        :param observation: Observation object
+        :type observation: AbstractObservation
+        """
+        self.obs: AbstractObservation = observation
+        self.current_observation: ObsType
+
+    def update(self, state: Dict) -> Dict:
+        """
+        Generate observation based on the current state of the simulation.
+
+        :param state: Simulation state dictionary
+        :type state: Dict
+        """
+        self.current_observation = self.obs.observe(state)
+        return self.current_observation
+
+    @property
+    def space(self) -> None:
+        """Gymnasium space object describing the observation space shape."""
+        return self.obs.space
+
+    @classmethod
+    def from_config(cls, config: Dict, game: "PrimaiteGame") -> "ObservationManager":
+        """Create observation space from a config.
+
+        :param config: Dictionary containing the configuration for this observation space.
+            It should contain the key 'type' which selects which observation class to use (from a choice of:
+            UC2BlueObservation, UC2RedObservation, UC2GreenObservation)
+            The other key is 'options' which are passed to the constructor of the selected observation class.
+        :type config: Dict
+        :param game: Reference to the PrimaiteGame object that spawned this observation.
+        :type game: PrimaiteGame
+        """
+        if config["type"] == "UC2BlueObservation":
+            return cls(UC2BlueObservation.from_config(config.get("options", {}), game=game))
+        elif config["type"] == "UC2RedObservation":
+            return cls(UC2RedObservation.from_config(config.get("options", {}), game=game))
+        elif config["type"] == "UC2GreenObservation":
+            return cls(UC2GreenObservation.from_config(config.get("options", {}), game=game))
+        else:
+            raise ValueError("Observation space type invalid")

src/primaite/game/agent/observations/observations.py

@@ -4,7 +4,6 @@ from ipaddress import IPv4Address
 from typing import Any, Dict, List, Optional, Tuple, TYPE_CHECKING
 
 from gymnasium import spaces
-from gymnasium.core import ObsType
 
 from primaite import getLogger
 from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
@@ -822,236 +821,3 @@ class ICSObservation(NullObservation):
     """ICS observation placeholder, currently not implemented so always returns a single 0."""
 
     pass
-
-
-class UC2BlueObservation(AbstractObservation):
-    """Container for all observations used by the blue agent in UC2.
-
-    TODO: there's no real need for a UC2 blue container class, we should be able to simply use the observation handler
-        for the purpose of compiling several observation components.
-    """
-
-    def __init__(
-        self,
-        nodes: List[NodeObservation],
-        links: List[LinkObservation],
-        acl: AclObservation,
-        ics: ICSObservation,
-        where: Optional[List[str]] = None,
-    ) -> None:
-        """Initialise UC2 blue observation.
-
-        :param nodes: List of node observations
-        :type nodes: List[NodeObservation]
-        :param links: List of link observations
-        :type links: List[LinkObservation]
-        :param acl: The Access Control List observation
-        :type acl: AclObservation
-        :param ics: The ICS observation
-        :type ics: ICSObservation
-        :param where: Where in the simulation state dict to find information. Not used in this particular observation
-            because it only compiles other observations and doesn't contribute any new information, defaults to None
-        :type where: Optional[List[str]], optional
-        """
-        super().__init__()
-        self.where: Optional[Tuple[str]] = where
-
-        self.nodes: List[NodeObservation] = nodes
-        self.links: List[LinkObservation] = links
-        self.acl: AclObservation = acl
-        self.ics: ICSObservation = ics
-
-        self.default_observation: Dict = {
-            "NODES": {i + 1: n.default_observation for i, n in enumerate(self.nodes)},
-            "LINKS": {i + 1: l.default_observation for i, l in enumerate(self.links)},
-            "ACL": self.acl.default_observation,
-            "ICS": self.ics.default_observation,
-        }
-
-    def observe(self, state: Dict) -> Dict:
-        """Generate observation based on the current state of the simulation.
-
-        :param state: Simulation state dictionary
-        :type state: Dict
-        :return: Observation
-        :rtype: Dict
-        """
-        if self.where is None:
-            return self.default_observation
-
-        obs = {}
-        obs["NODES"] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)}
-        obs["LINKS"] = {i + 1: link.observe(state) for i, link in enumerate(self.links)}
-        obs["ACL"] = self.acl.observe(state)
-        obs["ICS"] = self.ics.observe(state)
-
-        return obs
-
-    @property
-    def space(self) -> spaces.Space:
-        """
-        Gymnasium space object describing the observation space shape.
-
-        :return: Space
-        :rtype: spaces.Space
-        """
-        return spaces.Dict(
-            {
-                "NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}),
-                "LINKS": spaces.Dict({i + 1: link.space for i, link in enumerate(self.links)}),
-                "ACL": self.acl.space,
-                "ICS": self.ics.space,
-            }
-        )
-
-    @classmethod
-    def from_config(cls, config: Dict, game: "PrimaiteGame") -> "UC2BlueObservation":
-        """Create UC2 blue observation from a config.
-
-        :param config: Dictionary containing the configuration for this UC2 blue observation. This includes the nodes,
-            links, ACL and ICS observations.
-        :type config: Dict
-        :param game: Reference to the PrimaiteGame object that spawned this observation.
-        :type game: PrimaiteGame
-        :return: Constructed UC2 blue observation
-        :rtype: UC2BlueObservation
-        """
-        node_configs = config["nodes"]
-
-        num_services_per_node = config["num_services_per_node"]
-        num_folders_per_node = config["num_folders_per_node"]
-        num_files_per_folder = config["num_files_per_folder"]
-        num_nics_per_node = config["num_nics_per_node"]
-        nodes = [
-            NodeObservation.from_config(
-                config=n,
-                game=game,
-                num_services_per_node=num_services_per_node,
-                num_folders_per_node=num_folders_per_node,
-                num_files_per_folder=num_files_per_folder,
-                num_nics_per_node=num_nics_per_node,
-            )
-            for n in node_configs
-        ]
-
-        link_configs = config["links"]
-        links = [LinkObservation.from_config(config=link, game=game) for link in link_configs]
-
-        acl_config = config["acl"]
-        acl = AclObservation.from_config(config=acl_config, game=game)
-
-        ics_config = config["ics"]
-        ics = ICSObservation.from_config(config=ics_config, game=game)
-        new = cls(nodes=nodes, links=links, acl=acl, ics=ics, where=["network"])
-        return new
-
-
-class UC2RedObservation(AbstractObservation):
-    """Container for all observations used by the red agent in UC2."""
-
-    def __init__(self, nodes: List[NodeObservation], where: Optional[List[str]] = None) -> None:
-        super().__init__()
-        self.where: Optional[List[str]] = where
-        self.nodes: List[NodeObservation] = nodes
-
-        self.default_observation: Dict = {
-            "NODES": {i + 1: n.default_observation for i, n in enumerate(self.nodes)},
-        }
-
-    def observe(self, state: Dict) -> Dict:
-        """Generate observation based on the current state of the simulation."""
-        if self.where is None:
-            return self.default_observation
-
-        obs = {}
-        obs["NODES"] = {i + 1: node.observe(state) for i, node in enumerate(self.nodes)}
-        return obs
-
-    @property
-    def space(self) -> spaces.Space:
-        """Gymnasium space object describing the observation space shape."""
-        return spaces.Dict(
-            {
-                "NODES": spaces.Dict({i + 1: node.space for i, node in enumerate(self.nodes)}),
-            }
-        )
-
-    @classmethod
-    def from_config(cls, config: Dict, game: "PrimaiteGame") -> "UC2RedObservation":
-        """
-        Create UC2 red observation from a config.
-
-        :param config: Dictionary containing the configuration for this UC2 red observation.
-        :type config: Dict
-        :param game: Reference to the PrimaiteGame object that spawned this observation.
-        :type game: PrimaiteGame
-        """
-        node_configs = config["nodes"]
-        nodes = [NodeObservation.from_config(config=cfg, game=game) for cfg in node_configs]
-        return cls(nodes=nodes, where=["network"])
-
-
-class UC2GreenObservation(NullObservation):
-    """Green agent observation. As the green agent's actions don't depend on the observation, this is empty."""
-
-    pass
-
-
-class ObservationManager:
-    """
-    Manage the observations of an Agent.
-
-    The observation space has the purpose of:
-      1. Reading the outputted state from the PrimAITE Simulation.
-      2. Selecting parts of the simulation state that are requested by the simulation config
-      3. Formatting this information so an agent can use it to make decisions.
-    """
-
-    # TODO: Dear code reader: This class currently doesn't do much except hold an observation object. It will be changed
-    # to have more of it's own behaviour, and it will replace UC2BlueObservation and UC2RedObservation during the next
-    # refactor.
-
-    def __init__(self, observation: AbstractObservation) -> None:
-        """Initialise observation space.
-
-        :param observation: Observation object
-        :type observation: AbstractObservation
-        """
-        self.obs: AbstractObservation = observation
-        self.current_observation: ObsType
-
-    def update(self, state: Dict) -> Dict:
-        """
-        Generate observation based on the current state of the simulation.
-
-        :param state: Simulation state dictionary
-        :type state: Dict
-        """
-        self.current_observation = self.obs.observe(state)
-        return self.current_observation
-
-    @property
-    def space(self) -> None:
-        """Gymnasium space object describing the observation space shape."""
-        return self.obs.space
-
-    @classmethod
-    def from_config(cls, config: Dict, game: "PrimaiteGame") -> "ObservationManager":
-        """Create observation space from a config.
-
-        :param config: Dictionary containing the configuration for this observation space.
-            It should contain the key 'type' which selects which observation class to use (from a choice of:
-            UC2BlueObservation, UC2RedObservation, UC2GreenObservation)
-            The other key is 'options' which are passed to the constructor of the selected observation class.
-        :type config: Dict
-        :param game: Reference to the PrimaiteGame object that spawned this observation.
-        :type game: PrimaiteGame
-        """
-        if config["type"] == "UC2BlueObservation":
-            return cls(UC2BlueObservation.from_config(config.get("options", {}), game=game))
-        elif config["type"] == "UC2RedObservation":
-            return cls(UC2RedObservation.from_config(config.get("options", {}), game=game))
-        elif config["type"] == "UC2GreenObservation":
-            return cls(UC2GreenObservation.from_config(config.get("options", {}), game=game))
-        else:
-            raise ValueError("Observation space type invalid")

src/primaite/game/agent/scripted_agents/probabilistic_agent.py

@@ -7,7 +7,7 @@ from gymnasium.core import ObsType
 
 from primaite.game.agent.actions import ActionManager
 from primaite.game.agent.interface import AbstractScriptedAgent
-from primaite.game.agent.observations import ObservationManager
+from primaite.game.agent.observations.observation_manager import ObservationManager
 from primaite.game.agent.rewards import RewardFunction
 
 

src/primaite/game/agent/scripted_agents/random_agent.py

@@ -0,0 +1,21 @@
+from typing import Dict, Tuple
+
+from gymnasium.core import ObsType
+
+from primaite.game.agent.interface import AbstractScriptedAgent
+
+
+class RandomAgent(AbstractScriptedAgent):
+    """Agent that ignores its observation and acts completely at random."""
+
+    def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
+        """Sample the action space randomly.
+
+        :param obs: Current observation for this agent, not used in RandomAgent
+        :type obs: ObsType
+        :param timestep: The current simulation timestep, not used in RandomAgent
+        :type timestep: int
+        :return: Action formatted in CAOS format
+        :rtype: Tuple[str, Dict]
+        """
+        return self.action_manager.get_action(self.action_manager.space.sample())

src/primaite/game/game.py

@@ -6,11 +6,11 @@ from pydantic import BaseModel, ConfigDict
 
 from primaite import getLogger
 from primaite.game.agent.actions import ActionManager
-from primaite.game.agent.data_manipulation_bot import DataManipulationAgent
 from primaite.game.agent.interface import AbstractAgent, AgentSettings, ProxyAgent
-from primaite.game.agent.observations import ObservationManager
+from primaite.game.agent.observations.observation_manager import ObservationManager
 from primaite.game.agent.rewards import RewardFunction
-from primaite.game.agent.scripted_agents import ProbabilisticAgent
+from primaite.game.agent.scripted_agents.data_manipulation_bot import DataManipulationAgent
+from primaite.game.agent.scripted_agents.probabilistic_agent import ProbabilisticAgent
 from primaite.simulator.network.hardware.base import NodeOperatingState
 from primaite.simulator.network.hardware.nodes.host.computer import Computer
 from primaite.simulator.network.hardware.nodes.host.host_node import NIC

tests/conftest.py

@@ -10,7 +10,8 @@ from _pytest.monkeypatch import MonkeyPatch
 from primaite import getLogger, PRIMAITE_PATHS
 from primaite.game.agent.actions import ActionManager
 from primaite.game.agent.interface import AbstractAgent
-from primaite.game.agent.observations import ICSObservation, ObservationManager
+from primaite.game.agent.observations.observation_manager import ObservationManager
+from primaite.game.agent.observations.observations import ICSObservation
 from primaite.game.agent.rewards import RewardFunction
 from primaite.game.game import PrimaiteGame
 from primaite.session.session import PrimaiteSession

tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py

@@ -5,8 +5,8 @@ from typing import Union
 import yaml
 
 from primaite.config.load import data_manipulation_config_path
-from primaite.game.agent.data_manipulation_bot import DataManipulationAgent
 from primaite.game.agent.interface import ProxyAgent, RandomAgent
+from primaite.game.agent.scripted_agents.data_manipulation_bot import DataManipulationAgent
 from primaite.game.game import APPLICATION_TYPES_MAPPING, PrimaiteGame, SERVICE_TYPES_MAPPING
 from primaite.simulator.network.container import Network
 from primaite.simulator.network.hardware.nodes.host.computer import Computer

tests/integration_tests/game_layer/test_actions.py

@@ -10,28 +10,14 @@
 # 4. Check that the simulation has changed in the way that I expect.
 # 5. Repeat for all actions.
 
-from typing import Dict, Tuple
+from typing import Tuple
 
 import pytest
 
-from primaite.game.agent.actions import ActionManager
+from primaite.game.agent.interface import ProxyAgent
-from primaite.game.agent.interface import AbstractAgent, ProxyAgent
-from primaite.game.agent.observations import ICSObservation, ObservationManager
-from primaite.game.agent.rewards import RewardFunction
 from primaite.game.game import PrimaiteGame
 from primaite.simulator.file_system.file_system_item_abc import FileSystemItemHealthStatus
-from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
-from primaite.simulator.network.hardware.nodes.host.computer import Computer
-from primaite.simulator.network.hardware.nodes.host.server import Server
-from primaite.simulator.network.hardware.nodes.network.router import ACLAction, Router
-from primaite.simulator.network.hardware.nodes.network.switch import Switch
-from primaite.simulator.network.transmission.network_layer import IPProtocol
-from primaite.simulator.network.transmission.transport_layer import Port
-from primaite.simulator.sim_container import Simulation
 from primaite.simulator.system.applications.web_browser import WebBrowser
-from primaite.simulator.system.services.dns.dns_client import DNSClient
-from primaite.simulator.system.services.dns.dns_server import DNSServer
-from primaite.simulator.system.services.web_server.web_server import WebServer
 from primaite.simulator.system.software import SoftwareHealthState
 
 

tests/integration_tests/game_layer/test_observations.py

@@ -1,6 +1,6 @@
 from gymnasium import spaces
 
-from primaite.game.agent.observations import FileObservation
+from primaite.game.agent.observations.observations import FileObservation
 from primaite.simulator.network.hardware.nodes.host.computer import Computer
 from primaite.simulator.sim_container import Simulation
 

tests/integration_tests/network/test_capture_nmne.py

@@ -1,4 +1,4 @@
-from primaite.game.agent.observations import NicObservation
+from primaite.game.agent.observations.observations import NicObservation
 from primaite.simulator.network.hardware.nodes.host.server import Server
 from primaite.simulator.network.nmne import set_nmne_config
 from primaite.simulator.sim_container import Simulation

tests/unit_tests/_primaite/_game/_agent/test_probabilistic_agent.py

@@ -1,7 +1,8 @@
 from primaite.game.agent.actions import ActionManager
-from primaite.game.agent.observations import ICSObservation, ObservationManager
+from primaite.game.agent.observations.observation_manager import ObservationManager
+from primaite.game.agent.observations.observations import ICSObservation
 from primaite.game.agent.rewards import RewardFunction
-from primaite.game.agent.scripted_agents import ProbabilisticAgent
+from primaite.game.agent.scripted_agents.probabilistic_agent import ProbabilisticAgent
 
 
 def test_probabilistic_agent():