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

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

Browse Source
This commit is contained in:
SunilSamra
2023-06-06 13:21:04 +01:00
parent 940013f9a6
commit d922d4d054
5 changed files with 52 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/primaite/environment/primaite_env.py
View File

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

40
tests/conftest.py
View File

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

4
tests/test_observation_space.py
View File

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

2
tests/test_reward.py
View File

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

49
tests/test_single_action_space.py
View File

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