From 9a2fb2a0846cdc8958d7cad03fdbfaa6277e31bc Mon Sep 17 00:00:00 2001 From: Czar Echavez Date: Fri, 13 Sep 2024 11:11:58 +0100 Subject: [PATCH] #2880: fix action shape for num_ports + test --- src/primaite/game/agent/actions.py | 4 ++-- .../game_layer/test_action_shapes.py | 21 +++++++++++++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) create mode 100644 tests/integration_tests/game_layer/test_action_shapes.py diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py index a299788e..c864f75f 100644 --- a/src/primaite/game/agent/actions.py +++ b/src/primaite/game/agent/actions.py @@ -877,7 +877,7 @@ class FirewallACLRemoveRuleAction(AbstractAction): """Action which removes a rule from a firewall port's ACL.""" def __init__(self, manager: "ActionManager", max_acl_rules: int, **kwargs) -> None: - """Init method for RouterACLRemoveRuleAction. + """Init method for FirewallACLRemoveRuleAction. :param manager: Reference to the ActionManager which created this action. :type manager: ActionManager @@ -1524,7 +1524,7 @@ class ActionManager: "num_nics": max_nics_per_node, "num_acl_rules": max_acl_rules, "num_protocols": len(self.protocols), - "num_ports": len(self.protocols), + "num_ports": len(self.ports), "num_ips": len(self.ip_address_list), "max_acl_rules": max_acl_rules, "max_nics_per_node": max_nics_per_node, diff --git a/tests/integration_tests/game_layer/test_action_shapes.py b/tests/integration_tests/game_layer/test_action_shapes.py new file mode 100644 index 00000000..48500d8f --- /dev/null +++ b/tests/integration_tests/game_layer/test_action_shapes.py @@ -0,0 +1,21 @@ +# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK +from typing import Tuple + +from primaite.game.agent.interface import ProxyAgent +from primaite.game.game import PrimaiteGame +from tests import TEST_ASSETS_ROOT + +FIREWALL_ACTIONS_NETWORK = TEST_ASSETS_ROOT / "configs/firewall_actions_network.yaml" + + +def test_router_acl_add_rule_action_shape(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]): + """Test to check ROUTER_ADD_ACL_RULE has the expected action shape.""" + game, agent = game_and_agent + + # assert that the shape of the actions is correct + router_acl_add_rule_action = agent.action_manager.actions.get("ROUTER_ACL_ADDRULE") + assert router_acl_add_rule_action.shape.get("source_ip_id") == len(agent.action_manager.ip_address_list) + assert router_acl_add_rule_action.shape.get("dest_ip_id") == len(agent.action_manager.ip_address_list) + assert router_acl_add_rule_action.shape.get("source_port_id") == len(agent.action_manager.ports) + assert router_acl_add_rule_action.shape.get("dest_port_id") == len(agent.action_manager.ports) + assert router_acl_add_rule_action.shape.get("protocol_id") == len(agent.action_manager.protocols)