Fix comments

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__":

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"

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",

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)