From d922d4d0543d1f1dd300c147067cd6495156cdda Mon Sep 17 00:00:00 2001 From: SunilSamra Date: Tue, 6 Jun 2023 13:21:04 +0100 Subject: [PATCH] 893 - returned config_values in conftest to move run_generic_set_actions into test_single_action_space.py --- src/primaite/environment/primaite_env.py | 2 +- tests/conftest.py | 40 +------------------ tests/test_observation_space.py | 4 +- tests/test_reward.py | 2 +- tests/test_single_action_space.py | 49 +++++++++++++++++++++++- 5 files changed, 52 insertions(+), 45 deletions(-) diff --git a/src/primaite/environment/primaite_env.py b/src/primaite/environment/primaite_env.py index 4facb7b2..1794504a 100644 --- a/src/primaite/environment/primaite_env.py +++ b/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 diff --git a/tests/conftest.py b/tests/conftest.py index 3a99bcf6..740f65b7 100644 --- a/tests/conftest.py +++ b/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() diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py index 6a187761..d6eaa3b7 100644 --- a/tests/test_observation_space.py +++ b/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", diff --git a/tests/test_reward.py b/tests/test_reward.py index 4925a434..c54ee32f 100644 --- a/tests/test_reward.py +++ b/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", diff --git a/tests/test_single_action_space.py b/tests/test_single_action_space.py index fda4c96c..5fc6cb7e 100644 --- a/tests/test_single_action_space.py +++ b/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