#2869 - Minor changes trying to make pytest happy

This commit is contained in:
Charlie Crane
2025-01-06 11:38:07 +00:00
parent 505eab6ed9
commit 66d309871f
11 changed files with 17 additions and 23 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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