893 - returned config_values in conftest to move run_generic_set_actions into test_single_action_space.py

src/primaite/environment/primaite_env.py

@@ -368,7 +368,7 @@ class Primaite(Env):
             self.step_count,
             self.config_values,
         )
-        print(f"    Step {self.step_count} Reward: {str(reward)}")
+        # print(f"    Step {self.step_count} Reward: {str(reward)}")
         self.total_reward += reward
         if self.step_count == self.episode_steps:
             self.average_reward = self.total_reward / self.step_count

tests/conftest.py

@@ -169,10 +169,8 @@ def _get_primaite_env_from_config(
 
     if env.config_values.agent_identifier == "GENERIC":
         run_generic(env, config_values)
-    elif env.config_values.agent_identifier == "GENERIC_TEST":
-        run_generic_set_actions(env, config_values)
 
-    return env
+    return env, config_values
 
 
 def run_generic(env, config_values):
@@ -200,39 +198,3 @@ def run_generic(env, config_values):
         # env.reset()
 
     # env.close()
-
-
-def run_generic_set_actions(env, config_values):
-    """Run against a generic agent with specified blue agent actions."""
-    # Reset the environment at the start of the episode
-    # env.reset()
-    for episode in range(0, config_values.num_episodes):
-        for step in range(0, config_values.num_steps):
-            # Send the observation space to the agent to get an action
-            # TEMP - random action for now
-            # action = env.blue_agent_action(obs)
-            action = 0
-            if step == 5:
-                # [1, 1, 2, 1, 1, 1]
-                # Creates an ACL rule
-                # Deny traffic from server_1 to node_1 on port FTP
-                action = 7
-            elif step == 7:
-                # [1, 1, 2, 0] Node Action
-                # Sets Node 1 Hardware State to OFF
-                # Does not resolve any service
-                action = 16
-            # Run the simulation step on the live environment
-            obs, reward, done, info = env.step(action)
-
-            # Break if done is True
-            if done:
-                break
-
-            # Introduce a delay between steps
-            time.sleep(config_values.time_delay / 1000)
-
-        # Reset the environment at the end of the episode
-        # env.reset()
-
-    # env.close()

tests/test_observation_space.py

@@ -6,7 +6,7 @@ from tests.conftest import _get_primaite_env_from_config
 
 def test_creating_env_with_box_obs():
     """Try creating env with box observation space."""
-    env = _get_primaite_env_from_config(
+    env, config_values = _get_primaite_env_from_config(
         main_config_path=TEST_CONFIG_ROOT / "one_node_states_on_off_main_config.yaml",
         lay_down_config_path=TEST_CONFIG_ROOT / "box_obs_space_laydown_config.yaml",
     )
@@ -21,7 +21,7 @@ def test_creating_env_with_box_obs():
 
 def test_creating_env_with_multidiscrete_obs():
     """Try creating env with MultiDiscrete observation space."""
-    env = _get_primaite_env_from_config(
+    env, config_values = _get_primaite_env_from_config(
         main_config_path=TEST_CONFIG_ROOT / "one_node_states_on_off_main_config.yaml",
         lay_down_config_path=TEST_CONFIG_ROOT
         / "multidiscrete_obs_space_laydown_config.yaml",

tests/test_reward.py

@@ -8,7 +8,7 @@ def test_rewards_are_being_penalised_at_each_step_function():
 
     When the initial state is OFF compared to reference state which is ON.
     """
-    env = _get_primaite_env_from_config(
+    env, config_values = _get_primaite_env_from_config(
         main_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",

tests/test_single_action_space.py

@@ -1,15 +1,57 @@
+import time
+
 from primaite.common.enums import HardwareState
 from tests import TEST_CONFIG_ROOT
 from tests.conftest import _get_primaite_env_from_config
 
 
+def run_generic_set_actions(env, config_values):
+    """Run against a generic agent with specified blue agent actions."""
+    # Reset the environment at the start of the episode
+    # env.reset()
+    for episode in range(0, config_values.num_episodes):
+        for step in range(0, config_values.num_steps):
+            # Send the observation space to the agent to get an action
+            # TEMP - random action for now
+            # action = env.blue_agent_action(obs)
+            action = 0
+            print("Episode:", episode, "\nStep:", step)
+            if step == 5:
+                # [1, 1, 2, 1, 1, 1]
+                # Creates an ACL rule
+                # Deny traffic from server_1 to node_1 on port FTP
+                action = 7
+            elif step == 7:
+                # [1, 1, 2, 0] Node Action
+                # Sets Node 1 Hardware State to OFF
+                # Does not resolve any service
+                action = 16
+            # Run the simulation step on the live environment
+            obs, reward, done, info = env.step(action)
+
+            # Break if done is True
+            if done:
+                break
+
+            # Introduce a delay between steps
+            time.sleep(config_values.time_delay / 1000)
+
+        # Reset the environment at the end of the episode
+        # env.reset()
+
+    # 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(
+    env, config_values = _get_primaite_env_from_config(
         main_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",
     )
+
+    run_generic_set_actions(env, config_values)
+
     # 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
@@ -33,12 +75,15 @@ def test_single_action_space_is_valid():
 
 def test_agent_is_executing_actions_from_both_spaces():
     """Test to ensure the blue agent is carrying out both kinds of operations (NODE & ACL)."""
-    env = _get_primaite_env_from_config(
+    env, config_values = _get_primaite_env_from_config(
         main_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_generic_set_actions(env, config_values)
+
     # 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