#2606: add agent actions + test

This commit is contained in:
Czar Echavez
2024-05-20 13:10:21 +01:00
parent 406d4e889b
commit 80b1988ab9
4 changed files with 141 additions and 12 deletions

View File

@@ -313,13 +313,35 @@ class NodeFolderRestoreAction(NodeFolderAbstractAction):
self.verb: str = "restore"
class NodeFolderFileCreateAction(NodeFolderAbstractAction):
class NodeFileCreateAction(AbstractAction):
"""Action which creates a new file in a given folder."""
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None:
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs)
self.verb: str = "create"
def form_request(self, node_id: int, folder_name: str, file_name: str) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
node_name = self.manager.get_node_name_by_idx(node_id)
if node_name is None or folder_name is None or file_name is None:
return ["do_nothing"]
return ["network", "node", node_name, "file_system", "create", "file", folder_name, file_name]
class NodeFolderCreateAction(AbstractAction):
"""Action which creates a new folder."""
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None:
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs)
self.verb: str = "create"
def form_request(self, node_id: int, folder_name: str) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
node_name = self.manager.get_node_name_by_idx(node_id)
if node_name is None or folder_name is None:
return ["do_nothing"]
return ["network", "node", node_name, "file_system", "create", "folder", folder_name]
class NodeFileAbstractAction(AbstractAction):
"""Abstract base class for file actions.
@@ -401,6 +423,21 @@ class NodeFileCorruptAction(NodeFileAbstractAction):
self.verb: str = "corrupt"
class NodeFileAccessAction(AbstractAction):
"""Action which increases a file's access count."""
def __init__(self, manager: "ActionManager", num_nodes: int, num_folders: int, **kwargs) -> None:
super().__init__(manager, num_nodes=num_nodes, num_folders=num_folders, **kwargs)
self.verb: str = "create"
def form_request(self, node_id: int, folder_name: str, file_name: str) -> List[str]:
"""Return the action formatted as a request which can be ingested by the PrimAITE simulation."""
node_name = self.manager.get_node_name_by_idx(node_id)
if node_name is None or folder_name is None or file_name is None:
return ["do_nothing"]
return ["network", "node", node_name, "file_system", "access", folder_name, file_name]
class NodeAbstractAction(AbstractAction):
"""
Abstract base class for node actions.
@@ -854,11 +891,14 @@ class ActionManager:
"NODE_APPLICATION_INSTALL": NodeApplicationInstallAction,
"NODE_APPLICATION_REMOVE": NodeApplicationRemoveAction,
"NODE_FILE_SCAN": NodeFileScanAction,
"NODE_FILE_CREATE": NodeFileCreateAction,
"NODE_FILE_CHECKHASH": NodeFileCheckhashAction,
"NODE_FILE_DELETE": NodeFileDeleteAction,
"NODE_FILE_REPAIR": NodeFileRepairAction,
"NODE_FILE_RESTORE": NodeFileRestoreAction,
"NODE_FILE_CORRUPT": NodeFileCorruptAction,
"NODE_FILE_ACCESS": NodeFileAccessAction,
"NODE_FOLDER_CREATE": NodeFolderCreateAction,
"NODE_FOLDER_SCAN": NodeFolderScanAction,
"NODE_FOLDER_CHECKHASH": NodeFolderCheckhashAction,
"NODE_FOLDER_REPAIR": NodeFolderRepairAction,

View File

@@ -1,7 +1,7 @@
from __future__ import annotations
from pathlib import Path
from typing import Dict, Optional
from typing import Any, Dict, List, Optional
from prettytable import MARKDOWN, PrettyTable
@@ -64,21 +64,35 @@ class FileSystem(SimComponent):
)
self._create_manager = RequestManager()
def _create_file_action(request: List[Any], context: Any) -> RequestResponse:
file = self.create_file(folder_name=request[0], file_name=request[1])
if not file:
return RequestResponse.from_bool(False)
return RequestResponse(
status="success",
data={
"file_name": file.name,
"folder_name": file.folder_name,
"file_type": file.file_type,
"file_size": file.size,
},
)
self._create_manager.add_request(
name="file",
request_type=RequestType(
func=lambda request, context: RequestResponse.from_bool(
bool(self.create_file(folder_name=request[0], file_name=request[1]))
)
),
request_type=RequestType(func=_create_file_action),
)
def _create_folder_action(request: List[Any], context: Any) -> RequestResponse:
folder = self.create_folder(folder_name=request[0])
if not folder:
return RequestResponse.from_bool(False)
return RequestResponse(status="success", data={"folder_name": folder.name})
self._create_manager.add_request(
name="folder",
request_type=RequestType(
func=lambda request, context: RequestResponse.from_bool(
bool(self.create_folder(folder_name=request[0]))
)
),
request_type=RequestType(func=_create_folder_action),
)
rm.add_request(
name="create",

View File

@@ -431,12 +431,15 @@ def game_and_agent():
{"type": "NODE_APPLICATION_FIX"},
{"type": "NODE_APPLICATION_INSTALL"},
{"type": "NODE_APPLICATION_REMOVE"},
{"type": "NODE_FILE_CREATE"},
{"type": "NODE_FILE_SCAN"},
{"type": "NODE_FILE_CHECKHASH"},
{"type": "NODE_FILE_DELETE"},
{"type": "NODE_FILE_REPAIR"},
{"type": "NODE_FILE_RESTORE"},
{"type": "NODE_FILE_CORRUPT"},
{"type": "NODE_FILE_ACCESS"},
{"type": "NODE_FOLDER_CREATE"},
{"type": "NODE_FOLDER_SCAN"},
{"type": "NODE_FOLDER_CHECKHASH"},
{"type": "NODE_FOLDER_REPAIR"},

View File

@@ -329,6 +329,78 @@ def test_node_file_delete_integration(game_and_agent: Tuple[PrimaiteGame, ProxyA
assert file.deleted
def test_node_file_create(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]):
"""Test that a file is created."""
game, agent = game_and_agent
client_1 = game.simulation.network.get_node_by_hostname("client_1") #
action = (
"NODE_FILE_CREATE",
{
"node_id": 0,
"folder_name": "test",
"file_name": "file.txt",
},
)
agent.store_action(action)
game.step()
assert client_1.file_system.get_file(folder_name="test", file_name="file.txt")
def test_node_file_access(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]):
"""Test that the file access increments."""
game, agent = game_and_agent
client_1 = game.simulation.network.get_node_by_hostname("client_1") #
action = (
"NODE_FILE_CREATE",
{
"node_id": 0,
"folder_name": "test",
"file_name": "file.txt",
},
)
agent.store_action(action)
game.step()
assert client_1.file_system.get_file(folder_name="test", file_name="file.txt").num_access == 0
action = (
"NODE_FILE_ACCESS",
{
"node_id": 0,
"folder_name": "test",
"file_name": "file.txt",
},
)
agent.store_action(action)
game.step()
assert client_1.file_system.get_file(folder_name="test", file_name="file.txt").num_access == 1
def test_node_folder_create(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]):
"""Test that a folder is created."""
game, agent = game_and_agent
client_1 = game.simulation.network.get_node_by_hostname("client_1") #
action = (
"NODE_FOLDER_CREATE",
{
"node_id": 0,
"folder_name": "test",
},
)
agent.store_action(action)
game.step()
assert client_1.file_system.get_folder(folder_name="test")
def test_network_router_port_disable_integration(game_and_agent: Tuple[PrimaiteGame, ProxyAgent]):
"""Test that the NetworkPortDisableAction can form a request and that it is accepted by the simulation."""
game, agent = game_and_agent