#2656 - Unit tests for new ActionPenalty reward component, testing yaml and some minor changes to the implementation. Need to update Documentation to detail how this is added

tests/integration_tests/game_layer/test_rewards.py

@@ -2,7 +2,7 @@
 import yaml
 
 from primaite.game.agent.interface import AgentHistoryItem
-from primaite.game.agent.rewards import GreenAdminDatabaseUnreachablePenalty, WebpageUnavailablePenalty
+from primaite.game.agent.rewards import ActionPenalty, GreenAdminDatabaseUnreachablePenalty, WebpageUnavailablePenalty
 from primaite.game.game import PrimaiteGame
 from primaite.session.environment import PrimaiteGymEnv
 from primaite.simulator.network.hardware.nodes.host.server import Server
@@ -119,3 +119,67 @@ def test_shared_reward():
         g2_reward = env.game.agents["client_2_green_user"].reward_function.current_reward
         blue_reward = env.game.agents["defender"].reward_function.current_reward
         assert blue_reward == g1_reward + g2_reward
+
+
+def test_action_penalty_loads_from_config():
+    """Test to ensure that action penalty is correctly loaded from config into PrimaiteGymEnv"""
+    CFG_PATH = TEST_ASSETS_ROOT / "configs/action_penalty.yaml"
+    with open(CFG_PATH, "r") as f:
+        cfg = yaml.safe_load(f)
+
+    env = PrimaiteGymEnv(env_config=cfg)
+
+    env.reset()
+
+    ActionPenalty_Value = env.game.agents["defender"].reward_function.reward_components[2][0].penalty
+    CFG_Penalty_Value = cfg["agents"][3]["reward_function"]["reward_components"][2]["options"]["penalty_value"]
+
+    assert ActionPenalty_Value == CFG_Penalty_Value
+
+
+def test_action_penalty(game_and_agent):
+    """Test that the action penalty is correctly applied when agent performs any action"""
+
+    # Create an ActionPenalty Reward
+    Penalty = ActionPenalty(agent_name="Test_Blue_Agent", penalty=-1.0)
+
+    game, _ = game_and_agent
+
+    server_1: Server = game.simulation.network.get_node_by_hostname("server_1")
+    server_1.software_manager.install(DatabaseService)
+    db_service = server_1.software_manager.software.get("DatabaseService")
+    db_service.start()
+
+    client_1 = game.simulation.network.get_node_by_hostname("client_1")
+    client_1.software_manager.install(DatabaseClient)
+    db_client: DatabaseClient = client_1.software_manager.software.get("DatabaseClient")
+    db_client.configure(server_ip_address=server_1.network_interface[1].ip_address)
+    db_client.run()
+
+    response = db_client.apply_request(
+        [
+            "execute",
+        ]
+    )
+
+    state = game.get_sim_state()
+
+    # Assert that penalty is applied if action isn't DONOTHING
+    reward_value = Penalty.calculate(
+        state,
+        last_action_response=AgentHistoryItem(
+            timestep=0, action="NODE_APPLICATION_EXECUTE", parameters={}, request=["execute"], response=response
+        ),
+    )
+
+    assert reward_value == -1.0
+
+    # Assert that no penalty applied for a DONOTHING action
+    reward_value = Penalty.calculate(
+        state,
+        last_action_response=AgentHistoryItem(
+            timestep=0, action="DONOTHING", parameters={}, request=["execute"], response=response
+        ),
+    )
+
+    assert reward_value == 0