From 66d309871f18060fe3821aab158a6ff1e7b08770 Mon Sep 17 00:00:00 2001 From: Charlie Crane Date: Mon, 6 Jan 2025 11:38:07 +0000 Subject: [PATCH] #2869 - Minor changes trying to make pytest happy --- docs/source/how_to_guides/extensible_agents.rst | 4 ++-- src/primaite/game/agent/interface.py | 6 +++--- src/primaite/game/agent/scripted_agents/__init__.py | 7 +------ .../game/agent/scripted_agents/data_manipulation_bot.py | 1 - src/primaite/game/agent/scripted_agents/random_agent.py | 4 ++-- src/primaite/game/game.py | 7 +++---- tests/conftest.py | 3 ++- .../software_installation_and_configuration.py | 2 +- .../game_layer/observations/test_nic_observations.py | 2 +- tests/integration_tests/game_layer/test_rewards.py | 2 +- .../_primaite/_game/_agent/test_sticky_rewards.py | 2 +- 11 files changed, 17 insertions(+), 23 deletions(-) diff --git a/docs/source/how_to_guides/extensible_agents.rst b/docs/source/how_to_guides/extensible_agents.rst index b7c17b83..b9e00b60 100644 --- a/docs/source/how_to_guides/extensible_agents.rst +++ b/docs/source/how_to_guides/extensible_agents.rst @@ -7,13 +7,13 @@ Extensible Agents ***************** -Agents defined within PrimAITE have been updated to allow for easier creation of new bespoke agents. +Agents defined within PrimAITE have been updated to allow for easier creation of new bespoke agents. Developing Agents for PrimAITE ============================== -Agents within PrimAITE, follow the shown inheritance structure below. +Agents within PrimAITE, follow the shown inheritance structure below. # TODO: Turn this into an inheritance diagram diff --git a/src/primaite/game/agent/interface.py b/src/primaite/game/agent/interface.py index b980d748..6c1c633a 100644 --- a/src/primaite/game/agent/interface.py +++ b/src/primaite/game/agent/interface.py @@ -51,9 +51,9 @@ class AbstractAgent(BaseModel): history: List[AgentHistoryItem] = [] config: "AbstractAgent.ConfigSchema" - action_manager: ActionManager - observation_manager: ObservationManager - reward_function: RewardFunction + action_manager: "ActionManager" + observation_manager: "ObservationManager" + reward_function: "RewardFunction" class ConfigSchema(BaseModel): """ diff --git a/src/primaite/game/agent/scripted_agents/__init__.py b/src/primaite/game/agent/scripted_agents/__init__.py index 59e8bf40..5a97d15b 100644 --- a/src/primaite/game/agent/scripted_agents/__init__.py +++ b/src/primaite/game/agent/scripted_agents/__init__.py @@ -1,11 +1,6 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK from primaite.game.agent import interface -from primaite.game.agent.scripted_agents import ( - abstract_tap, - data_manipulation_bot, - probabilistic_agent, - random_agent, -) +from primaite.game.agent.scripted_agents import abstract_tap, data_manipulation_bot, probabilistic_agent, random_agent __all__ = ("abstract_tap", "data_manipulation_bot", "interface", "probabilistic_agent", "random_agent") diff --git a/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py b/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py index 3d9a7101..7ec119cf 100644 --- a/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py +++ b/src/primaite/game/agent/scripted_agents/data_manipulation_bot.py @@ -1,7 +1,6 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK from typing import Dict, Optional, Tuple - from gymnasium.core import ObsType from primaite.game.agent.scripted_agents.abstract_tap import AbstractTAPAgent diff --git a/src/primaite/game/agent/scripted_agents/random_agent.py b/src/primaite/game/agent/scripted_agents/random_agent.py index 89b19ece..14f642ef 100644 --- a/src/primaite/game/agent/scripted_agents/random_agent.py +++ b/src/primaite/game/agent/scripted_agents/random_agent.py @@ -37,6 +37,7 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="Periodic_Agent"): class ConfigSchema(AbstractScriptedAgent.ConfigSchema): """Configuration Schema for Periodic Agent.""" + agent_name: str = "Periodic_Agent" """Name of the agent.""" @@ -57,13 +58,12 @@ class PeriodicAgent(AbstractScriptedAgent, identifier="Periodic_Agent"): def start_variance(self) -> int: """Returns the deviation around the start step.""" return self.config.agent_settings.start_settings.variance - + @property def frequency(self) -> int: """Returns the number of timesteps to wait between performing actions.""" return self.config.agent_settings.start_settings.frequency - def _set_next_execution_timestep(self, timestep: int, variance: int) -> None: """Set the next execution timestep with a configured random variance. diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py index e83f59a6..db23eb14 100644 --- a/src/primaite/game/game.py +++ b/src/primaite/game/game.py @@ -525,23 +525,22 @@ class PrimaiteGame: agents_cfg = cfg.get("agents", []) for agent_cfg in agents_cfg: - agent_ref = agent_cfg["ref"] # noqa: F841 + agent_name = agent_cfg["ref"] # noqa: F841 agent_type = agent_cfg["type"] action_space_cfg = agent_cfg["action_space"] observation_space_cfg = agent_cfg["observation_space"] reward_function_cfg = agent_cfg["reward_function"] agent_settings = agent_cfg["agent_settings"] - # CREATE AGENT agent_config = { - "agent_name": agent_ref, + "agent_name": agent_name, "action_manager": action_space_cfg, "observation_manager": observation_space_cfg, "reward_function": reward_function_cfg, "agent_settings": agent_settings, } - # new_agent_cfg.update{} + # CREATE AGENT if agent_type in AbstractAgent._registry: new_agent = AbstractAgent._registry[agent_cfg["type"]].from_config(config=agent_config) # If blue agent is created, add to game.rl_agents diff --git a/tests/conftest.py b/tests/conftest.py index 28563333..0c211f49 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,9 +7,9 @@ from ray import init as rayinit 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.observation_manager import NestedObservation, ObservationManager from primaite.game.agent.rewards import RewardFunction -from primaite.game.agent.interface import AbstractAgent from primaite.game.agent.scripted_agents.probabilistic_agent import ProbabilisticAgent from primaite.game.game import PrimaiteGame from primaite.simulator.file_system.file_system import FileSystem @@ -369,6 +369,7 @@ def install_stuff_to_sim(sim: Simulation): # 5: Assert that the simulation starts off in the state that we expect assert len(sim.network.nodes) == 6 assert len(sim.network.links) == 5 + # 5.1: Assert the router is correctly configured r = sim.network.router_nodes[0] for i, acl_rule in enumerate(r.acl.acl): diff --git a/tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py b/tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py index 560fc44c..0ff6754d 100644 --- a/tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py +++ b/tests/integration_tests/configuration_file_parsing/software_installation_and_configuration.py @@ -6,8 +6,8 @@ from typing import Union import yaml from primaite.config.load import data_manipulation_config_path -from primaite.game.agent.scripted_agents.data_manipulation_bot import DataManipulationAgent from primaite.game.agent.interface import ProxyAgent +from primaite.game.agent.scripted_agents.data_manipulation_bot import DataManipulationAgent from primaite.game.agent.scripted_agents.probabilistic_agent import ProbabilisticAgent from primaite.game.game import PrimaiteGame, SERVICE_TYPES_MAPPING from primaite.simulator.network.container import Network diff --git a/tests/integration_tests/game_layer/observations/test_nic_observations.py b/tests/integration_tests/game_layer/observations/test_nic_observations.py index 6bcf6b4d..0ad03198 100644 --- a/tests/integration_tests/game_layer/observations/test_nic_observations.py +++ b/tests/integration_tests/game_layer/observations/test_nic_observations.py @@ -6,8 +6,8 @@ import pytest import yaml from gymnasium import spaces -from primaite.game.agent.observations.nic_observations import NICObservation from primaite.game.agent.interface import ProxyAgent +from primaite.game.agent.observations.nic_observations import NICObservation from primaite.game.game import PrimaiteGame from primaite.simulator.network.hardware.base import NetworkInterface from primaite.simulator.network.hardware.nodes.host.computer import Computer diff --git a/tests/integration_tests/game_layer/test_rewards.py b/tests/integration_tests/game_layer/test_rewards.py index c7aa6c99..dc7ed132 100644 --- a/tests/integration_tests/game_layer/test_rewards.py +++ b/tests/integration_tests/game_layer/test_rewards.py @@ -2,8 +2,8 @@ import pytest import yaml -from primaite.game.agent.rewards import ActionPenalty, GreenAdminDatabaseUnreachablePenalty, WebpageUnavailablePenalty from primaite.game.agent.interface import AgentHistoryItem +from primaite.game.agent.rewards import ActionPenalty, GreenAdminDatabaseUnreachablePenalty, WebpageUnavailablePenalty from primaite.game.game import PrimaiteGame from primaite.interface.request import RequestResponse from primaite.session.environment import PrimaiteGymEnv diff --git a/tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py b/tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py index de3f8144..0e4bf1bb 100644 --- a/tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py +++ b/tests/unit_tests/_primaite/_game/_agent/test_sticky_rewards.py @@ -1,11 +1,11 @@ # © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK +from primaite.game.agent.interface import AgentHistoryItem from primaite.game.agent.rewards import ( GreenAdminDatabaseUnreachablePenalty, WebpageUnavailablePenalty, WebServer404Penalty, ) -from primaite.game.agent.interface import AgentHistoryItem from primaite.interface.request import RequestResponse