diff --git a/agent/nebula_agent.py b/agent/nebula_agent.py index 026fb4a..6a6318f 100644 --- a/agent/nebula_agent.py +++ b/agent/nebula_agent.py @@ -1,7 +1,9 @@ """Nebula agent: responsible for persisting all types of messages.""" +import base64 import json import logging +import os from datetime import datetime from typing import Any @@ -19,14 +21,13 @@ ) logger = logging.getLogger(__name__) -CONFIG_HOME = "/root/.ostorlab" +SUPPORTED_FILE_TYPES = ["json"] class CustomEncoder(json.JSONEncoder): def default(self, obj: Any) -> Any: if isinstance(obj, bytes) is True: - return obj.decode("utf-8") - + return base64.b64encode(obj).decode("utf-8") return json.JSONEncoder.default(self, obj) @@ -40,6 +41,16 @@ def __init__( ) -> None: super().__init__(agent_definition, agent_settings) self._file_type = self.args.get("file_type", "json") + if self._file_type.lower() not in SUPPORTED_FILE_TYPES: + raise ValueError( + f"File type {self._file_type} is not supported. Supported file types are {SUPPORTED_FILE_TYPES}" + ) + + self._output_folder = ( + f"/output/messages_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}" + ) + if os.path.exists(self._output_folder) is False: + os.makedirs(self._output_folder) def process(self, message: m.Message) -> None: """Process the message and persist it to the file type and location specified in the agent definition. @@ -48,27 +59,22 @@ def process(self, message: m.Message) -> None: message: The message to process. """ logger.info("Processing message of selector : %s", message.selector) - message_data: dict[str, Any] | None = message.data - if message_data is None: - logger.warning("Message data is empty") - return None if self._file_type == "json": - self._persist_to_json(message_data) + self._persist_to_json(message) - def _persist_to_json(self, message_data: dict[str, Any]) -> None: - """Persist the message data to a JSON file. + def _persist_to_json(self, message_to_persist: m.Message) -> None: + """Persist message to JSON file. Args: - message_data: The message data to persist. + message_to_persist: The message to persist. """ - with open( - f"{CONFIG_HOME}/messages_{datetime.now().strftime('%Y-%m-%d_%H-%M')}.json", - "a", - ) as file: - file.write(json.dumps(message_data, cls=CustomEncoder) + "\n") + data = message_to_persist.data + selector = message_to_persist.selector + file_name = f"{self._output_folder}/{selector}_messages.json" - logger.info("Message persisted") + with open(file_name, "a") as file: + file.write(json.dumps(data, cls=CustomEncoder) + "\n") if __name__ == "__main__": diff --git a/ostorlab.yaml b/ostorlab.yaml index ff35051..b5b7647 100644 --- a/ostorlab.yaml +++ b/ostorlab.yaml @@ -67,12 +67,13 @@ description: | license: Apache-2.0 source: https://github.com/Ostorlab/agent_nebula -in_selectors: [] +in_selectors: + - v3 out_selectors: [] docker_file_path : Dockerfile docker_build_root : . mounts: - - '$CONFIG_HOME/:/root/.ostorlab/' + - '$CONFIG_HOME/:/output/' args: - name: "file_type" type: "string" diff --git a/tests/conftest.py b/tests/conftest.py index 2bdf385..9206f05 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -31,11 +31,6 @@ def fixture_agent( type="string", value=json.dumps("json").encode(), ), - utils_definitions.Arg( - name="file_path", - type="string", - value=json.dumps("output.json").encode(), - ), ], healthcheck_port=5301, redis_url="redis://guest:guest@localhost:6379", diff --git a/tests/nebula_agent_test.py b/tests/nebula_agent_test.py index 3b2a4ff..2cddd6e 100644 --- a/tests/nebula_agent_test.py +++ b/tests/nebula_agent_test.py @@ -1,46 +1,33 @@ """Unit tests for Nebula agent.""" -from ostorlab.agent.message import message as msg -from pytest_mock import plugin +import json +import pathlib -from agent import nebula_agent - - -def testAgentNebula_whenFileTypeIsJson_shouldPersistMessageToJSONFile( - nebula_test_agent: nebula_agent.NebulaAgent, - link_message: msg.Message, - mocker: plugin.MockerFixture, -) -> None: - """Test Nebula Agent when file_type is json and single message, - should persist message to JSON file.""" - open_mocker = mocker.patch("builtins.open", mocker.mock_open()) - - nebula_test_agent.process(link_message) - - assert open_mocker.called is True - path, _ = open_mocker.call_args[0] - assert nebula_agent.CONFIG_HOME + "/messages_" in path - assert any( - "https://ostorlab.co" in call_arg - for call_arg in open_mocker().write.call_args[0] - ) +from ostorlab.agent import definitions as agent_definitions +from ostorlab.runtimes import definitions as runtime_definitions +from ostorlab.utils import defintions as utils_definitions - -def testAgentNebula_whenMultipleMessages_shouldAppendMessagesToJSONFile( - nebula_test_agent: nebula_agent.NebulaAgent, - multiple_link_messages: list[msg.Message], - mocker: plugin.MockerFixture, -) -> None: - """Test Nebula Agent when multiple messages, should append messages to JSON file.""" - open_mocker = mocker.patch("builtins.open", mocker.mock_open()) - - for message in multiple_link_messages: - nebula_test_agent.process(message) - - assert open_mocker.called is True - assert open_mocker().write.call_count == len(multiple_link_messages) - assert all( - f"https://www.domain{i}.com" in call_arg - for i, args_list in enumerate(open_mocker().write.call_args_list) - for call_arg in args_list[0] - ) +from agent import nebula_agent +import pytest + + +def testAgentNebula_whenUnsupportedFileType_raisesValueError() -> None: + """Test that NebulaAgent raises ValueError when file type is not supported.""" + with pytest.raises(ValueError): + with (pathlib.Path(__file__).parent.parent / "ostorlab.yaml").open() as yaml_o: + definition = agent_definitions.AgentDefinition.from_yaml(yaml_o) + settings = runtime_definitions.AgentSettings( + key="agent/ostorlab/nebula", + bus_url="NA", + bus_exchange_topic="NA", + args=[ + utils_definitions.Arg( + name="file_type", + type="string", + value=json.dumps("txt").encode(), + ), + ], + healthcheck_port=5301, + redis_url="redis://guest:guest@localhost:6379", + ) + nebula_agent.NebulaAgent(definition, settings)