#2869 - Some additional test updates to amend failures. Pre-commit tbd and some cleanup

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

@@ -9,6 +9,7 @@ from gymnasium.core import ObsType
 
 from primaite.game.agent.scripted_agents.interface import AbstractAgent, AbstractScriptedAgent
 
+__all__ = ("AbstractTAPAgent")
 
 class AbstractTAPAgent(AbstractScriptedAgent, identifier="Abstract_TAP"):
     """Base class for TAP agents to inherit from."""
@@ -22,11 +23,6 @@ class AbstractTAPAgent(AbstractScriptedAgent, identifier="Abstract_TAP"):
 
         starting_node_name: Optional[str] = None
 
-    # @property
-    # def next_execution_timestep(self) -> int:
-    #     """Returns the agents next execution timestep."""
-    #     return self.next_execution_timestep
-
     @abstractmethod
     def get_action(self, obs: ObsType, timestep: int = 0) -> Tuple[str, Dict]:
         """Return an action to be taken in the environment."""

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

@@ -5,6 +5,7 @@ from gymnasium.core import ObsType
 
 from primaite.game.agent.scripted_agents.abstract_tap import AbstractTAPAgent
 
+__all__ = ("DataManipulationAgent")
 
 class DataManipulationAgent(AbstractTAPAgent, identifier="RedDatabaseCorruptingAgent"):
     """Agent that uses a DataManipulationBot to perform an SQL injection attack."""
@@ -17,11 +18,6 @@ class DataManipulationAgent(AbstractTAPAgent, identifier="RedDatabaseCorruptingA
 
         starting_application_name: Optional[str] = None
 
-    # def __init__(self, **kwargs: Any) -> None:
-    #     """Initialise DataManipulationAgent."""
-    #     self.setup_agent()
-    #     super().__init_subclass__(**kwargs)
-
     @property
     def starting_node_name(self) -> str:
         """Returns the agents starting node name."""

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

@@ -17,6 +17,11 @@ from primaite.interface.request import RequestFormat, RequestResponse
 if TYPE_CHECKING:
     pass
 
+__all__ = ("AgentHistoryItem",
+           "AgentStartSettings",
+           "AbstractAgent",
+           "AbstractScriptedAgent",
+           "ProxyAgent")
 
 class AgentHistoryItem(BaseModel):
     """One entry of an agent's action log - what the agent did and how the simulator responded in 1 step."""
@@ -116,7 +121,7 @@ class AbstractAgent(BaseModel):
         """
 
         model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
-        agent_name: Optional[str] = "Abstract_Agent"
+        agent_name: ClassVar[str] = "Abstract_Agent" # TODO: Make this a ClassVar[str] like verb in actions?
         history: List[AgentHistoryItem] = []
         _logger: AgentLog = AgentLog(agent_name=agent_name)
         action_manager: ActionManager

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

@@ -8,6 +8,7 @@ from gymnasium.core import ObsType
 
 from primaite.game.agent.scripted_agents.interface import AbstractScriptedAgent, AgentSettings
 
+__all__ = ("ProbabilisticAgent")
 
 class ProbabilisticAgent(AbstractScriptedAgent, identifier="ProbabilisticAgent"):
     """Scripted agent which randomly samples its action space with prescribed probabilities for each action."""

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

@@ -6,6 +6,7 @@ from gymnasium.core import ObsType
 
 from primaite.game.agent.scripted_agents.interface import AbstractScriptedAgent
 
+__all__ = ("RandomAgent", "PeriodicAgent")
 
 class RandomAgent(AbstractScriptedAgent, identifier="Random_Agent"):
     """Agent that ignores its observation and acts completely at random."""
@@ -38,18 +39,25 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="Periodic_Agent"):
 
     agent_name: str = "Periodic_Agent"
     """Name of the agent."""
+
+    # TODO: This is available in config.agent_settings.start_settings.start_step
     start_step: int = 20
     "The timestep at which an agent begins performing it's actions."
     start_variance: int = 5
     "Deviation around the start step."
+    
+    # TODO: This is available in config.agent_settings.start_settings.frequency
     frequency: int = 5
     "The number of timesteps to wait between performing actions."
+    
+    # TODO: This is available in config.agent_settings.start_settings.variance
     variance: int = 0
     "The amount the frequency can randomly change to."
     max_executions: int = 999999
     "Maximum number of times the agent can execute its action."
     num_executions: int = 0
     """Number of times the agent has executed an action."""
+    #TODO: Also in abstract_tap - move up and inherit? Add to AgentStartSettings?
     next_execution_timestep: int = 0
     """Timestep of the next action execution by the agent."""
 

tests/assets/configs/basic_switched_network.yaml

@@ -63,6 +63,10 @@ agents:
         start_step: 5
         frequency: 4
         variance: 3
+      action_probabilities:
+        0: 0.6
+        1: 0.4
+
 
 
 

tests/assets/configs/fix_duration_one_item.yaml

@@ -60,7 +60,9 @@ agents:
         start_step: 5
         frequency: 4
         variance: 3
-
+      action_probabilities:
+        0: 0.4
+        1: 0.6
 
 
   - ref: defender

tests/assets/configs/software_fix_duration.yaml

@@ -60,7 +60,9 @@ agents:
         start_step: 5
         frequency: 4
         variance: 3
-
+      action_probabilities:
+        0: 0.4
+        1: 0.6
 
 
   - ref: defender