oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'dev' into feature/1386-enable-a-repeatable-or-deterministic-baseline-test

Browse Source
This commit is contained in:
Czar Echavez
2023-07-04 09:41:07 +01:00
parent 409743dee9 6006f022a1
commit c7de7bf21b
71 changed files with 3963 additions and 2162 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
tests/config/legacy/legacy_training_config.yaml → tests/config/legacy_conversion/legacy_training_config.yaml
View File

21
tests/config/legacy/new_training_config.yaml → tests/config/legacy_conversion/new_training_config.yaml
View File

@@ -1,11 +1,20 @@
# Main Config File
# Generic config values
# Choose one of these (dependent on Agent being trained)
# "STABLE_BASELINES3_PPO"
# "STABLE_BASELINES3_A2C"
# "GENERIC"
agent_identifier: STABLE_BASELINES3_A2C
# Sets which agent algorithm framework will be used:
# "SB3" (Stable Baselines3)
# "RLLIB" (Ray[RLlib])
# "NONE" (Custom Agent)
agent_framework: SB3
# Sets which Red Agent algo/class will be used:
# "PPO" (Proximal Policy Optimization)
# "A2C" (Advantage Actor Critic)
# "HARDCODED" (Custom Agent)
# "RANDOM" (Random Action)
agent_identifier: PPO
# Sets How the Action Space is defined:
# "NODE"
# "ACL"
@@ -18,7 +27,7 @@ num_steps: 256
# Time delay between steps (for generic agents)
time_delay: 10
# Type of session to be run (TRAINING or EVALUATION)
session_type: TRAINING
session_type: TRAIN
# Determine whether to load an agent from file
load_agent: False
# File path and file name of agent if you're loading one in

27
tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml
View File

@@ -1,11 +1,22 @@
# Main Config File
# Training Config File
# Sets which agent algorithm framework will be used.
# Options are:
# "SB3" (Stable Baselines3)
# "RLLIB" (Ray RLlib)
# "CUSTOM" (Custom Agent)
agent_framework: SB3
# Sets which Agent class will be used.
# Options are:
# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework)
# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework)
# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type)
# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type)
# "RANDOM" (primaite.agents.simple.RandomAgent)
# "DUMMY" (primaite.agents.simple.DummyAgent)
agent_identifier: A2C
# Generic config values
# Choose one of these (dependent on Agent being trained)
# "STABLE_BASELINES3_PPO"
# "STABLE_BASELINES3_A2C"
# "GENERIC"
agent_identifier: STABLE_BASELINES3_A2C
# Sets How the Action Space is defined:
# "NODE"
# "ACL"
@@ -28,7 +39,7 @@ observation_space:
time_delay: 1
# Type of session to be run (TRAINING or EVALUATION)
session_type: TRAINING
session_type: TRAIN
# Determine whether to load an agent from file
load_agent: False
# File path and file name of agent if you're loading one in

27
tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml
View File

@@ -1,11 +1,22 @@
# Main Config File
# Training Config File
# Sets which agent algorithm framework will be used.
# Options are:
# "SB3" (Stable Baselines3)
# "RLLIB" (Ray RLlib)
# "CUSTOM" (Custom Agent)
agent_framework: CUSTOM
# Sets which Agent class will be used.
# Options are:
# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework)
# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework)
# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type)
# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type)
# "RANDOM" (primaite.agents.simple.RandomAgent)
# "DUMMY" (primaite.agents.simple.DummyAgent)
agent_identifier: RANDOM
# Generic config values
# Choose one of these (dependent on Agent being trained)
# "STABLE_BASELINES3_PPO"
# "STABLE_BASELINES3_A2C"
# "GENERIC"
agent_identifier: NONE
# Sets How the Action Space is defined:
# "NODE"
# "ACL"
@@ -24,7 +35,7 @@ observation_space:
time_delay: 1
# Filename of the scenario / laydown
session_type: TRAINING
session_type: TRAIN
# Determine whether to load an agent from file
load_agent: False
# File path and file name of agent if you're loading one in

27
tests/config/obs_tests/main_config_NODE_STATUSES.yaml
View File

@@ -1,11 +1,22 @@
# Main Config File
# Training Config File
# Sets which agent algorithm framework will be used.
# Options are:
# "SB3" (Stable Baselines3)
# "RLLIB" (Ray RLlib)
# "CUSTOM" (Custom Agent)
agent_framework: CUSTOM
# Sets which Agent class will be used.
# Options are:
# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework)
# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework)
# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type)
# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type)
# "RANDOM" (primaite.agents.simple.RandomAgent)
# "DUMMY" (primaite.agents.simple.DummyAgent)
agent_identifier: RANDOM
# Generic config values
# Choose one of these (dependent on Agent being trained)
# "STABLE_BASELINES3_PPO"
# "STABLE_BASELINES3_A2C"
# "GENERIC"
agent_identifier: NONE
# Sets How the Action Space is defined:
# "NODE"
# "ACL"
@@ -25,7 +36,7 @@ observation_space:
time_delay: 1
# Type of session to be run (TRAINING or EVALUATION)
session_type: TRAINING
session_type: TRAIN
# Determine whether to load an agent from file
load_agent: False
# File path and file name of agent if you're loading one in

27
tests/config/obs_tests/main_config_without_obs.yaml
View File

@@ -1,11 +1,22 @@
# Main Config File
# Training Config File
# Sets which agent algorithm framework will be used.
# Options are:
# "SB3" (Stable Baselines3)
# "RLLIB" (Ray RLlib)
# "CUSTOM" (Custom Agent)
agent_framework: CUSTOM
# Sets which Agent class will be used.
# Options are:
# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework)
# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework)
# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type)
# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type)
# "RANDOM" (primaite.agents.simple.RandomAgent)
# "DUMMY" (primaite.agents.simple.DummyAgent)
agent_identifier: RANDOM
# Generic config values
# Choose one of these (dependent on Agent being trained)
# "STABLE_BASELINES3_PPO"
# "STABLE_BASELINES3_A2C"
# "GENERIC"
agent_identifier: NONE
# Sets How the Action Space is defined:
# "NODE"
# "ACL"
@@ -18,7 +29,7 @@ num_steps: 5
# Time delay between steps (for generic agents)
time_delay: 1
# Type of session to be run (TRAINING or EVALUATION)
session_type: TRAINING
session_type: TRAIN
# Determine whether to load an agent from file
load_agent: False
# File path and file name of agent if you're loading one in

26
tests/config/one_node_states_on_off_main_config.yaml
View File

@@ -1,10 +1,22 @@
# Main Config File
# Training Config File
# Sets which agent algorithm framework will be used.
# Options are:
# "SB3" (Stable Baselines3)
# "RLLIB" (Ray RLlib)
# "CUSTOM" (Custom Agent)
agent_framework: CUSTOM
# Sets which Agent class will be used.
# Options are:
# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework)
# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework)
# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type)
# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type)
# "RANDOM" (primaite.agents.simple.RandomAgent)
# "DUMMY" (primaite.agents.simple.DummyAgent)
agent_identifier: DUMMY
# Generic config values
# Choose one of these (dependent on Agent being trained)
# "STABLE_BASELINES3_PPO"
# "STABLE_BASELINES3_A2C"
agent_identifier: GENERIC
# Sets How the Action Space is defined:
# "NODE"
# "ACL"
@@ -18,7 +30,7 @@ num_steps: 15
time_delay: 1
# Type of session to be run (TRAINING or EVALUATION)
session_type: TRAINING
session_type: EVAL
# Determine whether to load an agent from file
load_agent: False
# File path and file name of agent if you're loading one in

27
tests/config/single_action_space_fixed_blue_actions_main_config.yaml
View File

@@ -1,11 +1,22 @@
# Main Config File
# Training Config File
# Sets which agent algorithm framework will be used.
# Options are:
# "SB3" (Stable Baselines3)
# "RLLIB" (Ray RLlib)
# "CUSTOM" (Custom Agent)
agent_framework: CUSTOM
# Sets which Agent class will be used.
# Options are:
# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework)
# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework)
# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type)
# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type)
# "RANDOM" (primaite.agents.simple.RandomAgent)
# "DUMMY" (primaite.agents.simple.DummyAgent)
agent_identifier: RANDOM
# Generic config values
# Choose one of these (dependent on Agent being trained)
# "STABLE_BASELINES3_PPO"
# "STABLE_BASELINES3_A2C"
# "GENERIC"
agent_identifier: GENERIC
# Sets How the Action Space is defined:
# "NODE"
# "ACL"
@@ -18,7 +29,7 @@ num_steps: 15
# Time delay between steps (for generic agents)
time_delay: 1
# Type of session to be run (TRAINING or EVALUATION)
session_type: TRAINING
session_type: EVAL
# Determine whether to load an agent from file
load_agent: False
# File path and file name of agent if you're loading one in

27
tests/config/single_action_space_main_config.yaml
View File

@@ -1,11 +1,22 @@
# Main Config File
# Training Config File
# Sets which agent algorithm framework will be used.
# Options are:
# "SB3" (Stable Baselines3)
# "RLLIB" (Ray RLlib)
# "CUSTOM" (Custom Agent)
agent_framework: CUSTOM
# Sets which Agent class will be used.
# Options are:
# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework)
# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework)
# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type)
# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type)
# "RANDOM" (primaite.agents.simple.RandomAgent)
# "DUMMY" (primaite.agents.simple.DummyAgent)
agent_identifier: RANDOM
# Generic config values
# Choose one of these (dependent on Agent being trained)
# "STABLE_BASELINES3_PPO"
# "STABLE_BASELINES3_A2C"
# "GENERIC"
agent_identifier: GENERIC
# Sets How the Action Space is defined:
# "NODE"
# "ACL"
@@ -18,7 +29,7 @@ num_steps: 5
# Time delay between steps (for generic agents)
time_delay: 1
# Type of session to be run (TRAINING or EVALUATION)
session_type: TRAINING
session_type: EVAL
# Determine whether to load an agent from file
load_agent: False
# File path and file name of agent if you're loading one in

112
tests/config/test_random_red_main_config.yaml Normal file
View File

@@ -0,0 +1,112 @@
# Training Config File
# Sets which agent algorithm framework will be used.
# Options are:
# "SB3" (Stable Baselines3)
# "RLLIB" (Ray RLlib)
# "CUSTOM" (Custom Agent)
agent_framework: CUSTOM
# Sets which Agent class will be used.
# Options are:
# "A2C" (Advantage Actor Critic coupled with either SB3 or RLLIB agent_framework)
# "PPO" (Proximal Policy Optimization coupled with either SB3 or RLLIB agent_framework)
# "HARDCODED" (The HardCoded agents coupled with an ACL or NODE action_type)
# "DO_NOTHING" (The DoNothing agents coupled with an ACL or NODE action_type)
# "RANDOM" (primaite.agents.simple.RandomAgent)
# "DUMMY" (primaite.agents.simple.DummyAgent)
agent_identifier: DUMMY
# Sets whether Red Agent POL and IER is randomised.
# Options are:
# True
# False
random_red_agent: True
# Sets How the Action Space is defined:
# "NODE"
# "ACL"
# "ANY" node and acl actions
action_type: NODE
# Number of episodes to run per session
num_episodes: 2
# Number of time_steps per episode
num_steps: 15
# Time delay between steps (for generic agents)
time_delay: 1
# Type of session to be run (TRAINING or EVALUATION)
session_type: EVAL
# Determine whether to load an agent from file
load_agent: False
# File path and file name of agent if you're loading one in
agent_load_file: C:\[Path]\[agent_saved_filename.zip]
# Environment config values
# The high value for the observation space
observation_space_high_value: 1000000000
# Reward values
# Generic
all_ok: 0
# Node Hardware State
off_should_be_on: -10
off_should_be_resetting: -5
on_should_be_off: -2
on_should_be_resetting: -5
resetting_should_be_on: -5
resetting_should_be_off: -2
resetting: -3
# Node Software or Service State
good_should_be_patching: 2
good_should_be_compromised: 5
good_should_be_overwhelmed: 5
patching_should_be_good: -5
patching_should_be_compromised: 2
patching_should_be_overwhelmed: 2
patching: -3
compromised_should_be_good: -20
compromised_should_be_patching: -20
compromised_should_be_overwhelmed: -20
compromised: -20
overwhelmed_should_be_good: -20
overwhelmed_should_be_patching: -20
overwhelmed_should_be_compromised: -20
overwhelmed: -20
# Node File System State
good_should_be_repairing: 2
good_should_be_restoring: 2
good_should_be_corrupt: 5
good_should_be_destroyed: 10
repairing_should_be_good: -5
repairing_should_be_restoring: 2
repairing_should_be_corrupt: 2
repairing_should_be_destroyed: 0
repairing: -3
restoring_should_be_good: -10
restoring_should_be_repairing: -2
restoring_should_be_corrupt: 1
restoring_should_be_destroyed: 2
restoring: -6
corrupt_should_be_good: -10
corrupt_should_be_repairing: -10
corrupt_should_be_restoring: -10
corrupt_should_be_destroyed: 2
corrupt: -10
destroyed_should_be_good: -20
destroyed_should_be_repairing: -20
destroyed_should_be_restoring: -20
destroyed_should_be_corrupt: -20
destroyed: -20
scanning: -2
# IER status
red_ier_running: -5
green_ier_blocked: -10
# Patching / Reset durations
os_patching_duration: 5 # The time taken to patch the OS
node_reset_duration: 5 # The time taken to reset a node (hardware)
service_patching_duration: 5 # The time taken to patch a service
file_system_repairing_limit: 5 # The time take to repair the file system
file_system_restoring_limit: 5 # The time take to restore the file system
file_system_scanning_limit: 5 # The time taken to scan the file system

158
tests/conftest.py
View File

@@ -1,46 +1,151 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
import datetime
import shutil
import tempfile
import time
from datetime import datetime
from pathlib import Path
from typing import Final, Union
from typing import Dict, Union
from unittest.mock import patch
import pandas as pd
import pytest
from primaite import getLogger
from primaite.common.enums import AgentIdentifier
from primaite.environment.primaite_env import Primaite
from primaite.main import _get_session_path, _write_session_metadata_file
from primaite.primaite_session import PrimaiteSession
from primaite.utils.session_output_reader import av_rewards_dict
from tests.mock_and_patch.get_session_path_mock import get_temp_session_path
ACTION_SPACE_NODE_VALUES = 1
ACTION_SPACE_NODE_ACTION_VALUES = 1
_LOGGER = getLogger(__name__)
def _get_temp_session_path(session_timestamp: datetime) -> Path:
class TempPrimaiteSession(PrimaiteSession):
"""
A temporary PrimaiteSession class.
Uses context manager for deletion of files upon exit.
"""
def __init__(
self,
training_config_path: Union[str, Path],
lay_down_config_path: Union[str, Path],
):
super().__init__(training_config_path, lay_down_config_path)
self.setup()
def learn_av_reward_per_episode(self) -> Dict[int, float]:
"""Get the learn av reward per episode from file."""
csv_file = f"average_reward_per_episode_{self.timestamp_str}.csv"
return av_rewards_dict(self.learning_path / csv_file)
def eval_av_reward_per_episode_csv(self) -> Dict[int, float]:
"""Get the eval av reward per episode from file."""
csv_file = f"average_reward_per_episode_{self.timestamp_str}.csv"
return av_rewards_dict(self.evaluation_path / csv_file)
@property
def env(self) -> Primaite:
"""Direct access to the env for ease of testing."""
return self._agent_session._env # noqa
def __enter__(self):
return self
def __exit__(self, type, value, tb):
shutil.rmtree(self.session_path)
shutil.rmtree(self.session_path.parent)
_LOGGER.debug(f"Deleted temp session directory: {self.session_path}")
@pytest.fixture
def temp_primaite_session(request):
"""
Provides a temporary PrimaiteSession instance.
It's temporary as it uses a temporary directory as the session path.
To use this fixture you need to:
- parametrize your test function with:
- "temp_primaite_session"
- [[path to training config, path to lay down config]]
- Include the temp_primaite_session fixture as a param in your test
function.
- use the temp_primaite_session as a context manager assigning is the
name 'session'.
.. code:: python
from primaite.config.lay_down_config import dos_very_basic_config_path
from primaite.config.training_config import main_training_config_path
@pytest.mark.parametrize(
"temp_primaite_session",
[
[main_training_config_path(), dos_very_basic_config_path()]
],
indirect=True
)
def test_primaite_session(temp_primaite_session):
with temp_primaite_session as session:
# Learning outputs are saved in session.learning_path
session.learn()
# Evaluation outputs are saved in session.evaluation_path
session.evaluate()
# To ensure that all files are written, you must call .close()
session.close()
# If you need to inspect any session outputs, it must be done
# inside the context manager
# Now that we've exited the context manager, the
# session.session_path directory and its contents are deleted
"""
training_config_path = request.param[0]
lay_down_config_path = request.param[1]
with patch("primaite.agents.agent.get_session_path", get_temp_session_path) as mck:
mck.session_timestamp = datetime.now()
return TempPrimaiteSession(training_config_path, lay_down_config_path)
@pytest.fixture
def temp_session_path() -> Path:
"""
Get a temp directory session path the test session will output to.
:param session_timestamp: This is the datetime that the session started.
:return: The session directory path.
"""
session_timestamp = datetime.now()
date_dir = session_timestamp.strftime("%Y-%m-%d")
session_dir = session_timestamp.strftime("%Y-%m-%d_%H-%M-%S")
session_path = Path(tempfile.gettempdir()) / "primaite" / date_dir / session_dir
session_path = session_timestamp.strftime("%Y-%m-%d_%H-%M-%S")
session_path = Path(tempfile.gettempdir()) / "primaite" / date_dir / session_path
session_path.mkdir(exist_ok=True, parents=True)
return session_path
def _get_primaite_env_from_config(
training_config_path: Union[str, Path], lay_down_config_path: Union[str, Path]
training_config_path: Union[str, Path],
lay_down_config_path: Union[str, Path],
temp_session_path,
):
"""Takes a config path and returns the created instance of Primaite."""
session_timestamp: datetime = datetime.now()
session_path = _get_temp_session_path(session_timestamp)
session_path = temp_session_path(session_timestamp)
timestamp_str = session_timestamp.strftime("%Y-%m-%d_%H-%M-%S")
env = Primaite(
training_config_path=training_config_path,
lay_down_config_path=lay_down_config_path,
transaction_list=[],
session_path=session_path,
timestamp_str=timestamp_str,
)
@@ -49,7 +154,7 @@ def _get_primaite_env_from_config(
# TOOD: This needs t be refactored to happen outside. Should be part of
# a main Session class.
if env.training_config.agent_identifier == "GENERIC":
if env.training_config.agent_identifier is AgentIdentifier.RANDOM:
run_generic(env, config_values)
return env
@@ -95,9 +200,7 @@ def compare_file_content(output_a_file_path: str, output_b_file_path: str):
# both files have the same content
return True
# both files have different content
print(
f"{output_a_file_path} and {output_b_file_path} has different contents"
)
print(f"{output_a_file_path} and {output_b_file_path} has different contents")
return False
@@ -116,32 +219,3 @@ def compare_transaction_file(output_a_file_path: str, output_b_file_path: str):
# if the comparison is empty, both files are the same i.e. True
return data_a.compare(data_b).empty
class TestSession:
"""Class that contains session values."""
def __init__(self, training_config_path, laydown_config_path):
self.session_timestamp: Final[datetime] = datetime.now()
self.session_dir = _get_session_path(self.session_timestamp)
self.timestamp_str = self.session_timestamp.strftime("%Y-%m-%d_%H-%M-%S")
self.transaction_list = []
print(f"The output directory for this session is: {self.session_dir}")
self.env = Primaite(
training_config_path=training_config_path,
lay_down_config_path=laydown_config_path,
transaction_list=self.transaction_list,
session_path=self.session_dir,
timestamp_str=self.timestamp_str,
)
print("Writing Session Metadata file...")
_write_session_metadata_file(
session_dir=self.session_dir,
uuid="test",
session_timestamp=self.session_timestamp,
env=self.env,
)

8
tests/e2e_integration_tests/test_primaite_main.py
View File

@@ -1,8 +0,0 @@
from primaite.config.lay_down_config import data_manipulation_config_path
from primaite.config.training_config import main_training_config_path
from primaite.main import run
def test_primaite_main_e2e():
"""Tests the primaite.main.run function end-to-end."""
run(main_training_config_path(), data_manipulation_config_path())

407
tests/e2e_integration_tests/test_session_repeatability.py
View File

@@ -1,404 +1,3 @@
import time
from primaite import getLogger
from primaite.config.lay_down_config import data_manipulation_config_path
from primaite.main import (
_get_session_path,
_update_session_metadata_file,
run_generic,
run_stable_baselines3_a2c,
run_stable_baselines3_ppo,
)
from primaite.transactions.transactions_to_file import write_transaction_to_file
from tests import TEST_CONFIG_ROOT
from tests.conftest import TestSession, compare_file_content, compare_transaction_file
_LOGGER = getLogger(__name__)
def test_generic_same_results():
"""Runs seeded and deterministic Generic Primaite sessions and checks that the results are the same."""
print("")
print("=======================")
print("Generic test run")
print("=======================")
print("")
# run session 1
session1 = TestSession(
TEST_CONFIG_ROOT / "e2e/generic_deterministic_seeded_training_config.yaml",
data_manipulation_config_path(),
)
config_values = session1.env.training_config
# Get the number of steps (which is stored in the child config file)
config_values.num_steps = session1.env.episode_steps
run_generic(env=session1.env, config_values=session1.env.training_config)
_update_session_metadata_file(session_dir=session1.session_dir, env=session1.env)
# run session 2
session2 = TestSession(
TEST_CONFIG_ROOT / "e2e/generic_deterministic_seeded_training_config.yaml",
data_manipulation_config_path(),
)
config_values = session2.env.training_config
# Get the number of steps (which is stored in the child config file)
config_values.num_steps = session2.env.episode_steps
run_generic(env=session2.env, config_values=session2.env.training_config)
_update_session_metadata_file(session_dir=session2.session_dir, env=session2.env)
# wait until the csv files have been closed
while (not session1.env.csv_file.closed) or (not session2.env.csv_file.closed):
time.sleep(1)
# check if both outputs are the same
assert (
compare_file_content(
session1.env.csv_file.name,
session2.env.csv_file.name,
)
is True
)
# deterministic run
deterministic = TestSession(
TEST_CONFIG_ROOT / "e2e/generic_deterministic_seeded_training_config.yaml",
data_manipulation_config_path(),
)
deterministic.env.training_config.deterministic = True
run_generic(env=deterministic.env, config_values=deterministic.env.training_config)
_update_session_metadata_file(
session_dir=deterministic.session_dir, env=deterministic.env
)
# check if both outputs are the same
assert (
compare_file_content(
deterministic.env.csv_file.name,
TEST_CONFIG_ROOT
/ "e2e/deterministic_test_outputs/deterministic_generic.csv",
)
is True
)
def test_ppo_same_results():
"""Runs seeded and deterministic PPO Primaite sessions and checks that the results are the same."""
print("")
print("=======================")
print("PPO test run")
print("=======================")
print("")
training_session = TestSession(
TEST_CONFIG_ROOT / "e2e/ppo_deterministic_seeded_training_config.yaml",
data_manipulation_config_path(),
)
# Train agent
training_session.env.training_config.session_type = "TRAINING"
config_values = training_session.env.training_config
# Get the number of steps (which is stored in the child config file)
config_values.num_steps = training_session.env.episode_steps
run_stable_baselines3_ppo(
env=training_session.env,
config_values=config_values,
session_path=training_session.session_dir,
timestamp_str=training_session.timestamp_str,
)
write_transaction_to_file(
transaction_list=training_session.transaction_list,
session_path=training_session.session_dir,
timestamp_str=training_session.timestamp_str,
)
_update_session_metadata_file(
session_dir=training_session.session_dir, env=training_session.env
)
# Evaluate Agent again
eval_session1 = TestSession(
TEST_CONFIG_ROOT / "e2e/ppo_deterministic_seeded_training_config.yaml",
data_manipulation_config_path(),
)
# Get the number of steps (which is stored in the child config file)
config_values.num_steps = eval_session1.env.episode_steps
eval_session1.env.training_config.session_type = "EVALUATE"
# load the agent that was trained previously
eval_session1.env.training_config.load_agent = True
eval_session1.env.training_config.agent_load_file = (
_get_session_path(training_session.session_timestamp)
/ f"agent_saved_{training_session.timestamp_str}.zip"
)
config_values = eval_session1.env.training_config
run_stable_baselines3_ppo(
env=eval_session1.env,
config_values=config_values,
session_path=eval_session1.session_dir,
timestamp_str=eval_session1.timestamp_str,
)
write_transaction_to_file(
transaction_list=eval_session1.transaction_list,
session_path=eval_session1.session_dir,
timestamp_str=eval_session1.timestamp_str,
)
_update_session_metadata_file(
session_dir=eval_session1.session_dir, env=eval_session1.env
)
eval_session2 = TestSession(
TEST_CONFIG_ROOT / "e2e/ppo_deterministic_seeded_training_config.yaml",
data_manipulation_config_path(),
)
# Get the number of steps (which is stored in the child config file)
config_values.num_steps = eval_session2.env.episode_steps
eval_session2.env.training_config.session_type = "EVALUATE"
# load the agent that was trained previously
eval_session2.env.training_config.load_agent = True
eval_session2.env.training_config.agent_load_file = (
_get_session_path(training_session.session_timestamp)
/ f"agent_saved_{training_session.timestamp_str}.zip"
)
config_values = eval_session2.env.training_config
run_stable_baselines3_ppo(
env=eval_session2.env,
config_values=config_values,
session_path=eval_session2.session_dir,
timestamp_str=eval_session2.timestamp_str,
)
write_transaction_to_file(
transaction_list=eval_session2.transaction_list,
session_path=eval_session2.session_dir,
timestamp_str=eval_session2.timestamp_str,
)
_update_session_metadata_file(
session_dir=eval_session2.session_dir, env=eval_session2.env
)
# check if both eval outputs are the same
assert (
compare_transaction_file(
eval_session1.session_dir
/ f"all_transactions_{eval_session1.timestamp_str}.csv",
eval_session2.session_dir
/ f"all_transactions_{eval_session2.timestamp_str}.csv",
)
is True
)
# deterministic run
deterministic = TestSession(
TEST_CONFIG_ROOT / "e2e/ppo_deterministic_seeded_training_config.yaml",
data_manipulation_config_path(),
)
deterministic.env.training_config.deterministic = True
run_stable_baselines3_ppo(
env=deterministic.env,
config_values=config_values,
session_path=deterministic.session_dir,
timestamp_str=deterministic.timestamp_str,
)
write_transaction_to_file(
transaction_list=deterministic.transaction_list,
session_path=deterministic.session_dir,
timestamp_str=deterministic.timestamp_str,
)
_update_session_metadata_file(
session_dir=deterministic.session_dir, env=deterministic.env
)
# check if both outputs are the same
assert (
compare_transaction_file(
deterministic.session_dir
/ f"all_transactions_{deterministic.timestamp_str}.csv",
TEST_CONFIG_ROOT / "e2e/deterministic_test_outputs/deterministic_ppo.csv",
)
is True
)
def test_a2c_same_results():
"""Runs seeded and deterministic A2C Primaite sessions and checks that the results are the same."""
print("")
print("=======================")
print("A2C test run")
print("=======================")
print("")
training_session = TestSession(
TEST_CONFIG_ROOT / "e2e/a2c_deterministic_seeded_training_config.yaml",
data_manipulation_config_path(),
)
# Train agent
training_session.env.training_config.session_type = "TRAINING"
config_values = training_session.env.training_config
# Get the number of steps (which is stored in the child config file)
config_values.num_steps = training_session.env.episode_steps
run_stable_baselines3_a2c(
env=training_session.env,
config_values=config_values,
session_path=training_session.session_dir,
timestamp_str=training_session.timestamp_str,
)
write_transaction_to_file(
transaction_list=training_session.transaction_list,
session_path=training_session.session_dir,
timestamp_str=training_session.timestamp_str,
)
_update_session_metadata_file(
session_dir=training_session.session_dir, env=training_session.env
)
# Evaluate Agent again
eval_session1 = TestSession(
TEST_CONFIG_ROOT / "e2e/a2c_deterministic_seeded_training_config.yaml",
data_manipulation_config_path(),
)
# Get the number of steps (which is stored in the child config file)
config_values.num_steps = eval_session1.env.episode_steps
eval_session1.env.training_config.session_type = "EVALUATE"
# load the agent that was trained previously
eval_session1.env.training_config.load_agent = True
eval_session1.env.training_config.agent_load_file = (
_get_session_path(training_session.session_timestamp)
/ f"agent_saved_{training_session.timestamp_str}.zip"
)
config_values = eval_session1.env.training_config
run_stable_baselines3_a2c(
env=eval_session1.env,
config_values=config_values,
session_path=eval_session1.session_dir,
timestamp_str=eval_session1.timestamp_str,
)
write_transaction_to_file(
transaction_list=eval_session1.transaction_list,
session_path=eval_session1.session_dir,
timestamp_str=eval_session1.timestamp_str,
)
_update_session_metadata_file(
session_dir=eval_session1.session_dir, env=eval_session1.env
)
eval_session2 = TestSession(
TEST_CONFIG_ROOT / "e2e/a2c_deterministic_seeded_training_config.yaml",
data_manipulation_config_path(),
)
# Get the number of steps (which is stored in the child config file)
config_values.num_steps = eval_session2.env.episode_steps
eval_session2.env.training_config.session_type = "EVALUATE"
# load the agent that was trained previously
eval_session2.env.training_config.load_agent = True
eval_session2.env.training_config.agent_load_file = (
_get_session_path(training_session.session_timestamp)
/ f"agent_saved_{training_session.timestamp_str}.zip"
)
config_values = eval_session2.env.training_config
run_stable_baselines3_a2c(
env=eval_session2.env,
config_values=config_values,
session_path=eval_session2.session_dir,
timestamp_str=eval_session2.timestamp_str,
)
write_transaction_to_file(
transaction_list=eval_session2.transaction_list,
session_path=eval_session2.session_dir,
timestamp_str=eval_session2.timestamp_str,
)
_update_session_metadata_file(
session_dir=eval_session2.session_dir, env=eval_session2.env
)
# check if both eval outputs are the same
assert (
compare_transaction_file(
eval_session1.session_dir
/ f"all_transactions_{eval_session1.timestamp_str}.csv",
eval_session2.session_dir
/ f"all_transactions_{eval_session2.timestamp_str}.csv",
)
is True
)
# deterministic run
deterministic = TestSession(
TEST_CONFIG_ROOT / "e2e/a2c_deterministic_seeded_training_config.yaml",
data_manipulation_config_path(),
)
deterministic.env.training_config.deterministic = True
run_stable_baselines3_a2c(
env=deterministic.env,
config_values=config_values,
session_path=deterministic.session_dir,
timestamp_str=deterministic.timestamp_str,
)
write_transaction_to_file(
transaction_list=deterministic.transaction_list,
session_path=deterministic.session_dir,
timestamp_str=deterministic.timestamp_str,
)
_update_session_metadata_file(
session_dir=deterministic.session_dir, env=deterministic.env
)
# check if both outputs are the same
assert (
compare_transaction_file(
deterministic.session_dir
/ f"all_transactions_{deterministic.timestamp_str}.csv",
TEST_CONFIG_ROOT / "e2e/deterministic_test_outputs/deterministic_a2c.csv",
)
is True
)
def temp():
"""TODO rewrite tests."""
pass

0
tests/e2e_integration_tests/__init__.py → tests/mock_and_patch/__init__.py
View File

22
tests/mock_and_patch/get_session_path_mock.py Normal file
View File

@@ -0,0 +1,22 @@
import tempfile
from datetime import datetime
from pathlib import Path
from primaite import getLogger
_LOGGER = getLogger(__name__)
def get_temp_session_path(session_timestamp: datetime) -> Path:
"""
Get a temp directory session path the test session will output to.
:param session_timestamp: This is the datetime that the session started.
:return: The session directory path.
"""
date_dir = session_timestamp.strftime("%Y-%m-%d")
session_path = session_timestamp.strftime("%Y-%m-%d_%H-%M-%S")
session_path = Path(tempfile.gettempdir()) / "primaite" / date_dir / session_path
session_path.mkdir(exist_ok=True, parents=True)
_LOGGER.debug(f"Created temp session directory: {session_path}")
return session_path

4
tests/test_acl.py
View File

@@ -95,8 +95,6 @@ def test_rule_hash():
rule = ACLRule("DENY", "192.168.1.1", "192.168.1.2", "TCP", "80")
hash_value_local = hash(rule)
hash_value_remote = acl.get_dictionary_hash(
"DENY", "192.168.1.1", "192.168.1.2", "TCP", "80"
)
hash_value_remote = acl.get_dictionary_hash("DENY", "192.168.1.1", "192.168.1.2", "TCP", "80")
assert hash_value_local == hash_value_remote

212
tests/test_observation_space.py
View File

@@ -2,84 +2,79 @@
import numpy as np
import pytest
from primaite.environment.observations import (
NodeLinkTable,
NodeStatuses,
ObservationsHandler,
)
from primaite.environment.primaite_env import Primaite
from primaite.environment.observations import NodeLinkTable, NodeStatuses, ObservationsHandler
from tests import TEST_CONFIG_ROOT
from tests.conftest import _get_primaite_env_from_config
@pytest.fixture
def env(request):
"""Build Primaite environment for integration tests of observation space."""
marker = request.node.get_closest_marker("env_config_paths")
training_config_path = marker.args[0]["training_config_path"]
lay_down_config_path = marker.args[0]["lay_down_config_path"]
env = _get_primaite_env_from_config(
training_config_path=training_config_path,
lay_down_config_path=lay_down_config_path,
)
yield env
@pytest.mark.env_config_paths(
dict(
training_config_path=TEST_CONFIG_ROOT
/ "obs_tests/main_config_without_obs.yaml",
lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
)
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "obs_tests/main_config_without_obs.yaml",
TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
]
],
indirect=True,
)
def test_default_obs_space(env: Primaite):
def test_default_obs_space(temp_primaite_session):
"""Create environment with no obs space defined in config and check that the default obs space was created."""
env.update_environent_obs()
with temp_primaite_session as session:
session.env.update_environent_obs()
components = env.obs_handler.registered_obs_components
components = session.env.obs_handler.registered_obs_components
assert len(components) == 1
assert isinstance(components[0], NodeLinkTable)
assert len(components) == 1
assert isinstance(components[0], NodeLinkTable)
@pytest.mark.env_config_paths(
dict(
training_config_path=TEST_CONFIG_ROOT
/ "obs_tests/main_config_without_obs.yaml",
lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
)
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "obs_tests/main_config_without_obs.yaml",
TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
]
],
indirect=True,
)
def test_registering_components(env: Primaite):
def test_registering_components(temp_primaite_session):
"""Test regitering and deregistering a component."""
handler = ObservationsHandler()
component = NodeStatuses(env)
handler.register(component)
assert component in handler.registered_obs_components
handler.deregister(component)
assert component not in handler.registered_obs_components
with temp_primaite_session as session:
env = session.env
handler = ObservationsHandler()
component = NodeStatuses(env)
handler.register(component)
assert component in handler.registered_obs_components
handler.deregister(component)
assert component not in handler.registered_obs_components
@pytest.mark.env_config_paths(
dict(
training_config_path=TEST_CONFIG_ROOT
/ "obs_tests/main_config_NODE_LINK_TABLE.yaml",
lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
)
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "obs_tests/main_config_NODE_LINK_TABLE.yaml",
TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
]
],
indirect=True,
)
class TestNodeLinkTable:
"""Test the NodeLinkTable observation component (in isolation)."""
def test_obs_shape(self, env: Primaite):
def test_obs_shape(self, temp_primaite_session):
"""Try creating env with box observation space."""
env.update_environent_obs()
with temp_primaite_session as session:
env = session.env
env.update_environent_obs()
# we have three nodes and two links, with two service
# therefore the box observation space will have:
# * 5 rows (3 nodes + 2 links)
# * 6 columns (four fixed and two for the services)
assert env.env_obs.shape == (5, 6)
# we have three nodes and two links, with two service
# therefore the box observation space will have:
# * 5 rows (3 nodes + 2 links)
# * 6 columns (four fixed and two for the services)
assert env.env_obs.shape == (5, 6)
def test_value(self, env: Primaite):
def test_value(self, temp_primaite_session):
"""Test that the observation is generated correctly.
The laydown has:
@@ -125,36 +120,43 @@ class TestNodeLinkTable:
* 999 (999 traffic service1)
* 0 (no traffic for service2)
"""
# act = np.asarray([0,])
obs, reward, done, info = env.step(0) # apply the 'do nothing' action
with temp_primaite_session as session:
env = session.env
# act = np.asarray([0,])
obs, reward, done, info = env.step(0) # apply the 'do nothing' action
assert np.array_equal(
obs,
[
[1, 1, 3, 1, 1, 1],
[2, 1, 1, 1, 1, 4],
[3, 1, 1, 1, 0, 0],
[4, 0, 0, 0, 999, 0],
[5, 0, 0, 0, 999, 0],
],
)
assert np.array_equal(
obs,
[
[1, 1, 3, 1, 1, 1],
[2, 1, 1, 1, 1, 4],
[3, 1, 1, 1, 0, 0],
[4, 0, 0, 0, 999, 0],
[5, 0, 0, 0, 999, 0],
],
)
@pytest.mark.env_config_paths(
dict(
training_config_path=TEST_CONFIG_ROOT
/ "obs_tests/main_config_NODE_STATUSES.yaml",
lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
)
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "obs_tests/main_config_NODE_STATUSES.yaml",
TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
]
],
indirect=True,
)
class TestNodeStatuses:
"""Test the NodeStatuses observation component (in isolation)."""
def test_obs_shape(self, env: Primaite):
def test_obs_shape(self, temp_primaite_session):
"""Try creating env with NodeStatuses as the only component."""
assert env.env_obs.shape == (15,)
with temp_primaite_session as session:
env = session.env
assert env.env_obs.shape == (15,)
def test_values(self, env: Primaite):
def test_values(self, temp_primaite_session):
"""Test that the hardware and software states are encoded correctly.
The laydown has:
@@ -181,28 +183,36 @@ class TestNodeStatuses:
* service 1 = n/a (0)
* service 2 = n/a (0)
"""
obs, _, _, _ = env.step(0) # apply the 'do nothing' action
assert np.array_equal(obs, [1, 3, 1, 1, 1, 1, 1, 1, 1, 4, 1, 1, 1, 0, 0])
with temp_primaite_session as session:
env = session.env
obs, _, _, _ = env.step(0) # apply the 'do nothing' action
print(obs)
assert np.array_equal(obs, [1, 3, 1, 1, 1, 1, 1, 1, 1, 4, 1, 1, 1, 0, 0])
@pytest.mark.env_config_paths(
dict(
training_config_path=TEST_CONFIG_ROOT
/ "obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml",
lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
)
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml",
TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
]
],
indirect=True,
)
class TestLinkTrafficLevels:
"""Test the LinkTrafficLevels observation component (in isolation)."""
def test_obs_shape(self, env: Primaite):
def test_obs_shape(self, temp_primaite_session):
"""Try creating env with MultiDiscrete observation space."""
env.update_environent_obs()
with temp_primaite_session as session:
env = session.env
env.update_environent_obs()
# we have two links and two services, so the shape should be 2 * 2
assert env.env_obs.shape == (2 * 2,)
# we have two links and two services, so the shape should be 2 * 2
assert env.env_obs.shape == (2 * 2,)
def test_values(self, env: Primaite):
def test_values(self, temp_primaite_session):
"""Test that traffic values are encoded correctly.
The laydown has:
@@ -212,12 +222,14 @@ class TestLinkTrafficLevels:
* an IER trying to send 999 bits of data over both links the whole time (via the first service)
* link bandwidth of 1000, therefore the utilisation is 99.9%
"""
obs, reward, done, info = env.step(0)
obs, reward, done, info = env.step(0)
with temp_primaite_session as session:
env = session.env
obs, reward, done, info = env.step(0)
obs, reward, done, info = env.step(0)
# the observation space has combine_service_traffic set to False, so the space has this format:
# [link1_service1, link1_service2, link2_service1, link2_service2]
# we send 999 bits of data via link1 and link2 on service 1.
# therefore the first and third elements should be 6 and all others 0
# (`7` corresponds to 100% utiilsation and `6` corresponds to 87.5%-100%)
assert np.array_equal(obs, [6, 0, 6, 0])
# the observation space has combine_service_traffic set to False, so the space has this format:
# [link1_service1, link1_service2, link2_service1, link2_service2]
# we send 999 bits of data via link1 and link2 on service 1.
# therefore the first and third elements should be 6 and all others 0
# (`7` corresponds to 100% utiilsation and `6` corresponds to 87.5%-100%)
assert np.array_equal(obs, [6, 0, 6, 0])

55
tests/test_primaite_session.py Normal file
View File

@@ -0,0 +1,55 @@
import os
import pytest
from primaite import getLogger
from primaite.config.lay_down_config import dos_very_basic_config_path
from primaite.config.training_config import main_training_config_path
_LOGGER = getLogger(__name__)
@pytest.mark.parametrize(
"temp_primaite_session",
[[main_training_config_path(), dos_very_basic_config_path()]],
indirect=True,
)
def test_primaite_session(temp_primaite_session):
"""Tests the PrimaiteSession class and its outputs."""
with temp_primaite_session as session:
session_path = session.session_path
assert session_path.exists()
session.learn()
# Learning outputs are saved in session.learning_path
session.evaluate()
# Evaluation outputs are saved in session.evaluation_path
# If you need to inspect any session outputs, it must be done inside
# the context manager
# Check that the metadata json file exists
assert (session_path / "session_metadata.json").exists()
# Check that the network png file exists
assert (session_path / f"network_{session.timestamp_str}.png").exists()
# Check that both the transactions and av reward csv files exist
for file in session.learning_path.iterdir():
if file.suffix == ".csv":
assert "all_transactions" in file.name or "average_reward_per_episode" in file.name
# Check that both the transactions and av reward csv files exist
for file in session.evaluation_path.iterdir():
if file.suffix == ".csv":
assert "all_transactions" in file.name or "average_reward_per_episode" in file.name
_LOGGER.debug("Inspecting files in temp session path...")
for dir_path, dir_names, file_names in os.walk(session_path):
for file in file_names:
path = os.path.join(dir_path, file)
file_str = path.split(str(session_path))[-1]
_LOGGER.debug(f" {file_str}")
# Now that we've exited the context manager, the session.session_path
# directory and its contents are deleted
assert not session_path.exists()

75
tests/test_red_random_agent_behaviour.py
View File

@@ -1,68 +1,30 @@
from datetime import datetime
import pytest
from primaite.config.lay_down_config import data_manipulation_config_path
from primaite.environment.primaite_env import Primaite
from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed
from tests import TEST_CONFIG_ROOT
from tests.conftest import _get_temp_session_path
def run_generic(env, config_values):
"""Run against a generic agent."""
# Reset the environment at the start of the episode
env.reset()
for episode in range(0, config_values.num_episodes):
for step in range(0, config_values.num_steps):
# Send the observation space to the agent to get an action
# TEMP - random action for now
# action = env.blue_agent_action(obs)
# action = env.action_space.sample()
action = 0
# Run the simulation step on the live environment
obs, reward, done, info = env.step(action)
# Break if done is True
if done:
break
# Reset the environment at the end of the episode
env.reset()
env.close()
def test_random_red_agent_behaviour():
"""
Test that hardware state is penalised at each step.
When the initial state is OFF compared to reference state which is ON.
"""
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "test_random_red_main_config.yaml",
data_manipulation_config_path(),
]
],
indirect=True,
)
def test_random_red_agent_behaviour(temp_primaite_session):
"""Test that red agent POL is randomised each episode."""
list_of_node_instructions = []
# RUN TWICE so we can make sure that red agent is randomised
for i in range(2):
"""Takes a config path and returns the created instance of Primaite."""
session_timestamp: datetime = datetime.now()
session_path = _get_temp_session_path(session_timestamp)
with temp_primaite_session as session:
session.evaluate()
list_of_node_instructions.append(session.env.red_node_pol)
timestamp_str = session_timestamp.strftime("%Y-%m-%d_%H-%M-%S")
env = Primaite(
training_config_path=TEST_CONFIG_ROOT
/ "one_node_states_on_off_main_config.yaml",
lay_down_config_path=data_manipulation_config_path(),
transaction_list=[],
session_path=session_path,
timestamp_str=timestamp_str,
)
# set red_agent_
env.training_config.random_red_agent = True
training_config = env.training_config
training_config.num_steps = env.episode_steps
run_generic(env, training_config)
# add red pol instructions to list
list_of_node_instructions.append(env.red_node_pol)
session.evaluate()
list_of_node_instructions.append(session.env.red_node_pol)
# compare instructions to make sure that red instructions are truly random
for index, instruction in enumerate(list_of_node_instructions):
@@ -73,5 +35,4 @@ def test_random_red_agent_behaviour():
print(f"{key} end step: {instruction.get_end_step()}")
print(f"{key} target node id: {instruction.get_target_node_id()}")
print("")
assert list_of_node_instructions[0].__ne__(list_of_node_instructions[1])

12
tests/test_resetting_node.py
View File

@@ -1,13 +1,7 @@
"""Used to test Active Node functions."""
import pytest
from primaite.common.enums import (
FileSystemState,
HardwareState,
NodeType,
Priority,
SoftwareState,
)
from primaite.common.enums import FileSystemState, HardwareState, NodeType, Priority, SoftwareState
from primaite.common.service import Service
from primaite.config.training_config import TrainingConfig
from primaite.nodes.active_node import ActiveNode
@@ -57,9 +51,7 @@ def test_node_boots_correctly(operating_state, expected_operating_state):
file_system_state="GOOD",
config_values=1,
)
service_attributes = Service(
name="node", port="80", software_state=SoftwareState.COMPROMISED
)
service_attributes = Service(name="node", port="80", software_state=SoftwareState.COMPROMISED)
service_node.add_service(service_attributes)
for x in range(5):

31
tests/test_reward.py
View File

@@ -1,21 +1,26 @@
import pytest
from tests import TEST_CONFIG_ROOT
from tests.conftest import _get_primaite_env_from_config
def test_rewards_are_being_penalised_at_each_step_function():
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "one_node_states_on_off_main_config.yaml",
TEST_CONFIG_ROOT / "one_node_states_on_off_lay_down_config.yaml",
]
],
indirect=True,
)
def test_rewards_are_being_penalised_at_each_step_function(
temp_primaite_session,
):
"""
Test that hardware state is penalised at each step.
When the initial state is OFF compared to reference state which is ON.
"""
env = _get_primaite_env_from_config(
training_config_path=TEST_CONFIG_ROOT
/ "one_node_states_on_off_main_config.yaml",
lay_down_config_path=TEST_CONFIG_ROOT
/ "one_node_states_on_off_lay_down_config.yaml",
)
"""
The config 'one_node_states_on_off_lay_down_config.yaml' has 15 steps:
On different steps, the laydown config has Pattern of Life (PoLs) which change a state of the node's attribute.
For example, turning the nodes' file system state to CORRUPT from its original state GOOD.
@@ -38,4 +43,8 @@ def test_rewards_are_being_penalised_at_each_step_function():
For the 4 steps where this occurs the average reward is:
Average Reward: -8 (-120 / 15)
"""
assert env.average_reward == -8.0
with temp_primaite_session as session:
session.evaluate()
session.close()
ev_rewards = session.eval_av_reward_per_episode_csv()
assert ev_rewards[1] == -8.0

120
tests/test_single_action_space.py
View File

@@ -1,9 +1,10 @@
import time
import pytest
from primaite.common.enums import HardwareState
from primaite.environment.primaite_env import Primaite
from tests import TEST_CONFIG_ROOT
from tests.conftest import _get_primaite_env_from_config
def run_generic_set_actions(env: Primaite):
@@ -17,7 +18,6 @@ def run_generic_set_actions(env: Primaite):
# TEMP - random action for now
# action = env.blue_agent_action(obs)
action = 0
print("Episode:", episode, "\nStep:", step)
if step == 5:
# [1, 1, 2, 1, 1, 1]
# Creates an ACL rule
@@ -44,59 +44,71 @@ def run_generic_set_actions(env: Primaite):
# env.close()
def test_single_action_space_is_valid():
"""Test to ensure the blue agent is using the ACL action space and is carrying out both kinds of operations."""
env = _get_primaite_env_from_config(
training_config_path=TEST_CONFIG_ROOT / "single_action_space_main_config.yaml",
lay_down_config_path=TEST_CONFIG_ROOT
/ "single_action_space_lay_down_config.yaml",
)
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "single_action_space_main_config.yaml",
TEST_CONFIG_ROOT / "single_action_space_lay_down_config.yaml",
]
],
indirect=True,
)
def test_single_action_space_is_valid(temp_primaite_session):
"""Test single action space is valid."""
with temp_primaite_session as session:
env = session.env
run_generic_set_actions(env)
# Retrieve the action space dictionary values from environment
env_action_space_dict = env.action_dict.values()
# Flags to check the conditions of the action space
contains_acl_actions = False
contains_node_actions = False
both_action_spaces = False
# Loop through each element of the list (which is every value from the dictionary)
for dict_item in env_action_space_dict:
# Node action detected
if len(dict_item) == 4:
contains_node_actions = True
# Link action detected
elif len(dict_item) == 6:
contains_acl_actions = True
# If both are there then the ANY action type is working
if contains_node_actions and contains_acl_actions:
both_action_spaces = True
# Check condition should be True
assert both_action_spaces
run_generic_set_actions(env)
# Retrieve the action space dictionary values from environment
env_action_space_dict = env.action_dict.values()
# Flags to check the conditions of the action space
contains_acl_actions = False
contains_node_actions = False
both_action_spaces = False
# Loop through each element of the list (which is every value from the dictionary)
for dict_item in env_action_space_dict:
# Node action detected
if len(dict_item) == 4:
contains_node_actions = True
# Link action detected
elif len(dict_item) == 6:
contains_acl_actions = True
# If both are there then the ANY action type is working
if contains_node_actions and contains_acl_actions:
both_action_spaces = True
# Check condition should be True
assert both_action_spaces
def test_agent_is_executing_actions_from_both_spaces():
@pytest.mark.parametrize(
"temp_primaite_session",
[
[
TEST_CONFIG_ROOT / "single_action_space_fixed_blue_actions_main_config.yaml",
TEST_CONFIG_ROOT / "single_action_space_lay_down_config.yaml",
]
],
indirect=True,
)
def test_agent_is_executing_actions_from_both_spaces(temp_primaite_session):
"""Test to ensure the blue agent is carrying out both kinds of operations (NODE & ACL)."""
env = _get_primaite_env_from_config(
training_config_path=TEST_CONFIG_ROOT
/ "single_action_space_fixed_blue_actions_main_config.yaml",
lay_down_config_path=TEST_CONFIG_ROOT
/ "single_action_space_lay_down_config.yaml",
)
# Run environment with specified fixed blue agent actions only
run_generic_set_actions(env)
# Retrieve hardware state of computer_1 node in laydown config
# Agent turned this off in Step 5
computer_node_hardware_state = env.nodes["1"].hardware_state
# Retrieve the Access Control List object stored by the environment at the end of the episode
access_control_list = env.acl
# Use the Access Control List object acl object attribute to get dictionary
# Use dictionary.values() to get total list of all items in the dictionary
acl_rules_list = access_control_list.acl.values()
# Length of this list tells you how many items are in the dictionary
# This number is the frequency of Access Control Rules in the environment
# In the scenario, we specified that the agent should create only 1 acl rule
num_of_rules = len(acl_rules_list)
# Therefore these statements below MUST be true
assert computer_node_hardware_state == HardwareState.OFF
assert num_of_rules == 1
with temp_primaite_session as session:
env = session.env
# Run environment with specified fixed blue agent actions only
run_generic_set_actions(env)
# Retrieve hardware state of computer_1 node in laydown config
# Agent turned this off in Step 5
computer_node_hardware_state = env.nodes["1"].hardware_state
# Retrieve the Access Control List object stored by the environment at the end of the episode
access_control_list = env.acl
# Use the Access Control List object acl object attribute to get dictionary
# Use dictionary.values() to get total list of all items in the dictionary
acl_rules_list = access_control_list.acl.values()
# Length of this list tells you how many items are in the dictionary
# This number is the frequency of Access Control Rules in the environment
# In the scenario, we specified that the agent should create only 1 acl rule
num_of_rules = len(acl_rules_list)
# Therefore these statements below MUST be true
assert computer_node_hardware_state == HardwareState.OFF
assert num_of_rules == 1

8
tests/test_training_config.py
View File

@@ -7,8 +7,8 @@ from tests import TEST_CONFIG_ROOT
def test_legacy_lay_down_config_yaml_conversion():
"""Tests the conversion of legacy lay down config files."""
legacy_path = TEST_CONFIG_ROOT / "legacy" / "legacy_training_config.yaml"
new_path = TEST_CONFIG_ROOT / "legacy" / "new_training_config.yaml"
legacy_path = TEST_CONFIG_ROOT / "legacy_conversion" / "legacy_training_config.yaml"
new_path = TEST_CONFIG_ROOT / "legacy_conversion" / "new_training_config.yaml"
with open(legacy_path, "r") as file:
legacy_dict = yaml.safe_load(file)
@@ -24,13 +24,13 @@ def test_legacy_lay_down_config_yaml_conversion():
def test_create_config_values_main_from_file():
"""Tests creating an instance of TrainingConfig from file."""
new_path = TEST_CONFIG_ROOT / "legacy" / "new_training_config.yaml"
new_path = TEST_CONFIG_ROOT / "legacy_conversion" / "new_training_config.yaml"
training_config.load(new_path)
def test_create_config_values_main_from_legacy_file():
"""Tests creating an instance of TrainingConfig from legacy file."""
new_path = TEST_CONFIG_ROOT / "legacy" / "legacy_training_config.yaml"
new_path = TEST_CONFIG_ROOT / "legacy_conversion" / "legacy_training_config.yaml"
training_config.load(new_path, legacy_file=True)