901 - merged with dev

tests/config/legacy_conversion/new_training_config.yaml

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

tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml

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

tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml

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

tests/config/obs_tests/main_config_NODE_STATUSES.yaml

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

tests/config/obs_tests/main_config_without_obs.yaml

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

tests/config/one_node_states_on_off_main_config.yaml

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

tests/config/single_action_space_fixed_blue_actions_main_config.yaml

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

tests/config/single_action_space_main_config.yaml

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

tests/config/test_random_red_main_config.yaml

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

tests/conftest.py

@@ -1,43 +1,150 @@
 # 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 Union
+from typing import Dict, Union
+from unittest.mock import patch
 
+import pytest
+
+from primaite import getLogger
+from primaite.common.enums import AgentIdentifier
 from primaite.environment.primaite_env import Primaite
+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,
     )
@@ -46,7 +153,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

tests/e2e_integration_tests/test_primaite_main.py

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

tests/mock_and_patch/get_session_path_mock.py

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

tests/test_acl.py

@@ -98,8 +98,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

tests/test_observation_space.py

@@ -5,17 +5,12 @@ import time
 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
 
 
-def run_generic_set_actions(env: Primaite):
+def run_generic_set_actions(env):
     """Run against a generic agent with specified blue agent actions."""
     # Reset the environment at the start of the episode
     env.reset()
@@ -70,61 +65,75 @@ def env(request):
     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:
@@ -170,36 +179,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, 0],
-                [2, 1, 1, 1, 1, 4, 1],
-                [3, 1, 1, 1, 0, 0, 2],
-                [4, 0, 0, 0, 999, 0, 3],
-                [5, 0, 0, 0, 999, 0, 4],
-            ],
-        )
+            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:
@@ -226,28 +242,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:
@@ -257,36 +281,42 @@ 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%)
-        print(obs)
-        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])
 
 
-@pytest.mark.env_config_paths(
-    dict(
-        training_config_path=TEST_CONFIG_ROOT
-        / "obs_tests/main_config_ACCESS_CONTROL_LIST.yaml",
-        lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
-    )
+@pytest.mark.parametrize(
+    "temp_primaite_session",
+    [
+        [
+            TEST_CONFIG_ROOT / "single_action_space_fixed_blue_actions_main_config.yaml",
+            TEST_CONFIG_ROOT / "obs_tests/laydown_ACL.yaml",
+        ]
+    ],
+    indirect=True,
 )
 class TestAccessControlList:
     """Test the AccessControlList 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 ACLs
         assert env.env_obs.shape == (6 * 3)
 
-    def test_values(self, env: Primaite):
+    def test_values(self, temp_primaite_session):
         """Test that traffic values are encoded correctly.
 
         The laydown has:
@@ -296,8 +326,10 @@ class TestAccessControlList:
             * 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]
@@ -307,16 +339,13 @@ class TestAccessControlList:
         print(obs)
         assert np.array_equal(obs, [])
 
-    def test_observation_space_with_implicit_rule(self):
+    def test_observation_space_with_implicit_rule(self, temp_primaite_session):
         """Test observation space is what is expected when an agent adds ACLs during an episode."""
         # Used to use env from test fixture but AtrributeError function object has no 'training_config'
-        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 / "obs_tests/laydown_ACL.yaml",
-        )
-        run_generic_set_actions(env)
-        obs = env.env_obs
+        with temp_primaite_session as session:
+            env = session.env
+            run_generic_set_actions(env)
+            obs = env.env_obs
         """
         Observation space at the end of the episode.
         At the start of the episode, there is a single implicit Deny rule = 1,1,1,1,1,0
@@ -327,6 +356,4 @@ class TestAccessControlList:
         """
 
         # assert current_obs == [2, 2, 3, 2, 3, 0,    2, 4, 2, 3, 3, 1,   1, 1, 1, 1, 1, 2]
-        assert np.array_equal(
-            obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2]
-        )
+        assert np.array_equal(obs, [2, 2, 3, 2, 3, 0, 2, 4, 2, 3, 3, 1, 1, 1, 1, 1, 1, 2])

tests/test_primaite_session.py

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

tests/test_red_random_agent_behaviour.py

@@ -0,0 +1,38 @@
+import pytest
+
+from primaite.config.lay_down_config import data_manipulation_config_path
+from primaite.nodes.node_state_instruction_red import NodeStateInstructionRed
+from tests import TEST_CONFIG_ROOT
+
+
+@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 = []
+
+    with temp_primaite_session as session:
+        session.evaluate()
+        list_of_node_instructions.append(session.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):
+        for key in list_of_node_instructions[index].keys():
+            instruction: NodeStateInstructionRed = list_of_node_instructions[index][key]
+            print(f"run {index}")
+            print(f"{key} start step: {instruction.get_start_step()}")
+            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])

tests/test_resetting_node.py

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

tests/test_reward.py

@@ -1,32 +1,50 @@
+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",
-    )
 
-    """
-    On different steps (of the 13 in total) these are the following rewards for config_6 which are activated:
-        File System State: goodShouldBeCorrupt = 5 (between Steps 1 & 3)
-        Hardware State: onShouldBeOff = -2 (between Steps 4 & 6)
-        Service State: goodShouldBeCompromised = 5 (between Steps 7 & 9)
-        Software State (Software State): goodShouldBeCompromised = 5 (between Steps 10 & 12)
+    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.
+        As a result these are the following rewards are activated:
+            File System State: corrupt_should_be_good = -10 * 2 (on Steps 1 & 2)
+            Hardware State: off_should_be_on = -10 * 2 (on Steps 4 & 5)
+            Service State: compromised_should_be_good = -20 * 2 (on Steps 7 & 8)
+            Software State: compromised_should_be_good = -20 * 2 (on Steps 10 & 11)
 
-    Total Reward: -2 - 2 + 5 + 5 + 5 + 5 + 5 + 5 = 26
-    Step Count: 13
+            The Pattern of Life (PoLs) last for 2 steps, so the agent is penalised twice.
+
+    Note: This test run inherits from conftest.py where the PrimAITE environment is ran and the blue agent is hard-coded
+    to do NOTHING on every step.
+    We use Pattern of Lifes (PoLs) to change the nodes states and display that the agent is being penalised on all steps
+    where the live network node differs from the network reference node.
+
+    Total Reward: -10 + -10 + -10 + -10 + -20 + -20 + -20 + -20 = -120
+    Step Count: 15
 
     For the 4 steps where this occurs the average reward is:
-        Average Reward: 2 (26 / 13)
+        Average Reward: -8 (-120 / 15)
     """
-    print("average reward", env.average_reward)
-    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

tests/test_single_action_space.py

@@ -1,15 +1,16 @@
 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):
     """Run against a generic agent with specified blue agent actions."""
     # Reset the environment at the start of the episode
-    env.reset()
+    # env.reset()
     training_config = env.training_config
     for episode in range(0, training_config.num_episodes):
         for step in range(0, training_config.num_steps):
@@ -44,63 +45,75 @@ 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) == 7:
-            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) == 7:
+                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
-    # 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
-    # This 1 rule added to the implicit deny means there should be 2 rules in total.
-    rules_count = 0
-    for rule in acl_rules_list:
-        if rule != -1:
-            rules_count += 1
-    # Therefore these statements below MUST be true
-    assert computer_node_hardware_state == HardwareState.OFF
-    assert rules_count == 2
+    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
+        # 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
+        # This 1 rule added to the implicit deny means there should be 2 rules in total.
+        rules_count = 0
+        for rule in acl_rules_list:
+            if rule != -1:
+                rules_count += 1
+        # Therefore these statements below MUST be true
+        assert computer_node_hardware_state == HardwareState.OFF
+        assert rules_count == 2

tests/test_training_config.py

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