update more tests to new actions schema

This commit is contained in:
Marek Wolan
2025-01-17 14:38:59 +00:00
parent 03cab0fcec
commit 858406c4a3
48 changed files with 351 additions and 413 deletions

View File

@@ -2,6 +2,7 @@
from unittest.mock import Mock
import pytest
from pydantic import ValidationError
from primaite.game.agent.actions import ActionManager
from primaite.game.agent.actions.manager import DoNothingAction
@@ -37,7 +38,7 @@ def test_do_nothing_action_form_request():
],
) # flake8: noqa
@pytest.mark.parametrize(
"node_name, service_name, expect_to_do_nothing",
"node_name, service_name, expect_failure",
[
("pc_1", "chrome", False),
(None, "chrome", True),
@@ -45,34 +46,15 @@ def test_do_nothing_action_form_request():
(None, None, True),
],
) # flake8: noqa
def test_service_action_form_request(node_name, service_name, expect_to_do_nothing, action_class, action_verb):
def test_service_action_form_request(node_name, service_name, expect_failure, action_class, action_verb):
"""Test that the ServiceScanAction can form a request and that it is correct."""
request = action_class.form_request(
config=action_class.ConfigSchema(node_name=node_name, service_name=service_name)
)
if expect_to_do_nothing:
assert request == ["do_nothing"]
if expect_failure:
with pytest.raises(ValidationError):
request = action_class.form_request(
config=action_class.ConfigSchema(node_name=node_name, service_name=service_name)
)
else:
request = action_class.form_request(
config=action_class.ConfigSchema(node_name=node_name, service_name=service_name)
)
assert request == ["network", "node", node_name, "service", service_name, action_verb]
@pytest.mark.parametrize(
"node_name, service_name, expect_to_do_nothing",
[
("pc_1", "chrome", False),
(None, "chrome", True),
("pc_1", None, True),
(None, None, True),
],
) # flake8: noqa
def test_service_scan_form_request(node_name, service_name, expect_to_do_nothing):
"""Test that the ServiceScanAction can form a request and that it is correct."""
request = NodeServiceScanAction.form_request(
NodeServiceScanAction.ConfigSchema(node_id=node_name, service_id=service_name)
)
if expect_to_do_nothing:
assert request == ["do_nothing"]
else:
assert request == ["network", "node", node_name, "service", service_name, "scan"]

View File

@@ -39,6 +39,8 @@ def test_creating_agent_from_dict():
}
agent = RandomAgent(
config={
"ref": "random_agent",
"team": "BLUE",
"action_space": action_config,
"observation_space": observation_config,
"reward_function": reward_config,

View File

@@ -98,7 +98,7 @@ class TestFileSystemRequiresScan:
"""
cfg = yaml.safe_load(obs_cfg_yaml)
manager = ObservationManager(cfg)
manager = ObservationManager(config=cfg)
hosts: List[HostObservation] = manager.obs.components["NODES"].hosts
for i, host in enumerate(hosts):

View File

@@ -92,7 +92,7 @@ class TestWebpageUnavailabilitySticky:
# agent did a successful fetch
action = "node_application_execute"
params = {"node_id": 0, "application_id": 0}
params = {"node_name": "computer", "application_name": "WebBrowser"}
request = ["network", "node", "computer", "application", "WebBrowser", "execute"]
response = RequestResponse(status="success", data={})
browser_history.append({"outcome": 200})
@@ -115,7 +115,7 @@ class TestWebpageUnavailabilitySticky:
# agent fails to fetch, get a -1.0 reward
action = "node_application_execute"
params = {"node_id": 0, "application_id": 0}
params = {"node_name": "computer", "application_name": "WebBrowser"}
request = ["network", "node", "computer", "application", "WebBrowser", "execute"]
response = RequestResponse(status="failure", data={})
browser_history.append({"outcome": 404})
@@ -127,7 +127,7 @@ class TestWebpageUnavailabilitySticky:
# agent fails again to fetch, get a -1.0 reward again
action = "node_application_execute"
params = {"node_id": 0, "application_id": 0}
params = {"node_name": "computer", "application_name": "WebBrowser"}
request = ["network", "node", "computer", "application", "WebBrowser", "execute"]
response = RequestResponse(status="failure", data={})
browser_history.append({"outcome": 404})
@@ -153,7 +153,7 @@ class TestWebpageUnavailabilitySticky:
# agent did a successful fetch
action = "node_application_execute"
params = {"node_id": 0, "application_id": 0}
params = {"node_name": "computer", "application_name": "WebBrowser"}
request = ["network", "node", "computer", "application", "WebBrowser", "execute"]
response = RequestResponse(status="success", data={})
browser_history.append({"outcome": 200})
@@ -175,7 +175,7 @@ class TestWebpageUnavailabilitySticky:
# agent fails to fetch, get a -1.0 reward
action = "node_application_execute"
params = {"node_id": 0, "application_id": 0}
params = {"node_name": "computer", "application_name": "WebBrowser"}
request = ["network", "node", "computer", "application", "WebBrowser", "execute"]
response = RequestResponse(status="failure", data={})
browser_history.append({"outcome": 404})
@@ -187,7 +187,7 @@ class TestWebpageUnavailabilitySticky:
# agent fails again to fetch, get a -1.0 reward again
action = "node_application_execute"
params = {"node_id": 0, "application_id": 0}
params = {"node_name": "computer", "application_name": "WebBrowser"}
request = ["network", "node", "computer", "application", "WebBrowser", "execute"]
response = RequestResponse(status="failure", data={})
browser_history.append({"outcome": 404})
@@ -217,7 +217,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
# agent did a successful fetch
action = "node_application_execute"
params = {"node_id": 0, "application_id": 0}
params = {"node_name": "computer", "application_name": "DatabaseClient"}
request = ["network", "node", "computer", "application", "DatabaseClient", "execute"]
response = RequestResponse(status="success", data={})
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
@@ -238,7 +238,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
# agent fails to fetch, get a -1.0 reward
action = "node_application_execute"
params = {"node_id": 0, "application_id": 0}
params = {"node_name": "computer", "application_name": "DatabaseClient"}
request = ["network", "node", "computer", "application", "DatabaseClient", "execute"]
response = RequestResponse(status="failure", data={})
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
@@ -249,7 +249,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
# agent fails again to fetch, get a -1.0 reward again
action = "node_application_execute"
params = {"node_id": 0, "application_id": 0}
params = {"node_name": "computer", "application_name": "DatabaseClient"}
request = ["network", "node", "computer", "application", "DatabaseClient", "execute"]
response = RequestResponse(status="failure", data={})
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
@@ -276,7 +276,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
# agent did a successful fetch
action = "node_application_execute"
params = {"node_id": 0, "application_id": 0}
params = {"node_name": "computer", "application_name": "DatabaseClient"}
request = ["network", "node", "computer", "application", "DatabaseClient", "execute"]
response = RequestResponse(status="success", data={})
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
@@ -297,7 +297,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
# agent fails to fetch, get a -1.0 reward
action = "node_application_execute"
params = {"node_id": 0, "application_id": 0}
params = {"node_name": "computer", "application_name": "DatabaseClient"}
request = ["network", "node", "computer", "application", "DatabaseClient", "execute"]
response = RequestResponse(status="failure", data={})
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}
@@ -308,7 +308,7 @@ class TestGreenAdminDatabaseUnreachableSticky:
# agent fails again to fetch, get a -1.0 reward again
action = "node_application_execute"
params = {"node_id": 0, "application_id": 0}
params = {"node_name": "computer", "application_name": "DatabaseClient"}
request = ["network", "node", "computer", "application", "DatabaseClient", "execute"]
response = RequestResponse(status="failure", data={})
state = {"network": {"nodes": {"computer": {"applications": {"DatabaseClient": {}}}}}}