#917 - Integrated the PrimaiteSession into all tests.

- Ran a full pre-commit hook and thus encountered tons of fixes required

tests/config/legacy/new_training_config.yaml

@@ -6,7 +6,7 @@
 # "SB3" (Stable Baselines3)
 # "RLLIB" (Ray[RLlib])
 # "NONE" (Custom Agent)
-agent_framework: RLLIB
+agent_framework: SB3
 
 # Sets which Red Agent algo/class will be used:
 # "PPO" (Proximal Policy Optimization)
@@ -27,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/conftest.py

@@ -1,37 +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 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):
+        del self._agent_session._env.episode_av_reward_writer
+        del self._agent_session._env.transaction_writer
+        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_path = session_timestamp.strftime("%Y-%m-%d_%H-%M-%S")
-    session_path = Path(tempfile.gettempdir()) / "primaite" / date_dir / session_path
+    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(
@@ -45,7 +159,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,24 @@
+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_active_node.py

@@ -60,7 +60,9 @@ def test_os_state_change_if_not_compromised(operating_state, expected_state):
         1,
     )
 
-    active_node.set_software_state_if_not_compromised(SoftwareState.OVERWHELMED)
+    active_node.set_software_state_if_not_compromised(
+        SoftwareState.OVERWHELMED
+    )
 
     assert active_node.software_state == expected_state
 
@@ -98,7 +100,9 @@ def test_file_system_change(operating_state, expected_state):
         (HardwareState.ON, FileSystemState.CORRUPT),
     ],
 )
-def test_file_system_change_if_not_compromised(operating_state, expected_state):
+def test_file_system_change_if_not_compromised(
+    operating_state, expected_state
+):
     """
     Test that a node cannot change its file system state.
 
@@ -116,6 +120,8 @@ def test_file_system_change_if_not_compromised(operating_state, expected_state):
         1,
     )
 
-    active_node.set_file_system_state_if_not_compromised(FileSystemState.CORRUPT)
+    active_node.set_file_system_state_if_not_compromised(
+        FileSystemState.CORRUPT
+    )
 
     assert active_node.file_system_state_actual == expected_state

tests/test_observation_space.py

@@ -7,79 +7,78 @@ from primaite.environment.observations import (
     NodeStatuses,
     ObservationsHandler,
 )
-from primaite.environment.primaite_env import Primaite
 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 +124,45 @@ 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 +189,38 @@ 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
+            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 +230,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])

tests/test_primaite_session.py

@@ -0,0 +1,61 @@
+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_resetting_node.py

@@ -18,7 +18,9 @@ from primaite.nodes.service_node import ServiceNode
     "starting_operating_state, expected_operating_state",
     [(HardwareState.RESETTING, HardwareState.ON)],
 )
-def test_node_resets_correctly(starting_operating_state, expected_operating_state):
+def test_node_resets_correctly(
+    starting_operating_state, expected_operating_state
+):
     """Tests that a node resets correctly."""
     active_node = ActiveNode(
         node_id="0",

tests/test_reward.py

@@ -1,26 +1,33 @@
+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:
+    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)
+        Software State (Software State): goodShouldBeCompromised = 5 (between
+        Steps 10 & 12)
 
     Total Reward: -2 - 2 + 5 + 5 + 5 + 5 + 5 + 5 = 26
     Step Count: 13
@@ -28,5 +35,8 @@ def test_rewards_are_being_penalised_at_each_step_function():
     For the 4 steps where this occurs the average reward is:
         Average Reward: 2 (26 / 13)
     """
-    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_service_node.py

@@ -45,7 +45,9 @@ def test_service_state_change(operating_state, expected_state):
         (HardwareState.ON, SoftwareState.OVERWHELMED),
     ],
 )
-def test_service_state_change_if_not_comprised(operating_state, expected_state):
+def test_service_state_change_if_not_comprised(
+    operating_state, expected_state
+):
     """
     Test that a node cannot change the state of a running service.
 
@@ -65,6 +67,8 @@ def test_service_state_change_if_not_comprised(operating_state, expected_state):
     service = Service("TCP", 80, SoftwareState.GOOD)
     service_node.add_service(service)
 
-    service_node.set_service_state_if_not_compromised("TCP", SoftwareState.OVERWHELMED)
+    service_node.set_service_state_if_not_compromised(
+        "TCP", SoftwareState.OVERWHELMED
+    )
 
     assert service_node.get_service_state("TCP") == expected_state

tests/test_single_action_space.py

@@ -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):
@@ -44,59 +45,72 @@ 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

tests/test_training_config.py

@@ -16,7 +16,9 @@ def test_legacy_lay_down_config_yaml_conversion():
     with open(new_path, "r") as file:
         new_dict = yaml.safe_load(file)
 
-    converted_dict = training_config.convert_legacy_training_config_dict(legacy_dict)
+    converted_dict = training_config.convert_legacy_training_config_dict(
+        legacy_dict
+    )
 
     for key, value in new_dict.items():
         assert converted_dict[key] == value