UT: Move and rename StubAgentLogRepository to InMemoryAgentLogRepository

monkey/tests/monkey_island/__init__.py

@@ -4,3 +4,4 @@ from .open_error_file_repository import OpenErrorFileRepository
 from .in_memory_agent_configuration_repository import InMemoryAgentConfigurationRepository
 from .in_memory_simulation_configuration import InMemorySimulationRepository
 from .in_memory_credentials_repository import InMemoryCredentialsRepository
+from .in_memory_agent_log_repository import InMemoryAgentLogRepository

monkey/tests/monkey_island/in_memory_agent_log_repository.py

@@ -0,0 +1,19 @@
+from common.types import AgentID
+from monkey_island.cc.repository import IAgentLogRepository, UnknownRecordError
+
+
+class InMemoryAgentLogRepository(IAgentLogRepository):
+    def __init__(self):
+        self._agent_logs = {}
+
+    def upsert_agent_log(self, agent_id: AgentID, log_contents: str):
+        if agent_id not in self._agent_logs.keys():
+            self._agent_logs[agent_id] = log_contents
+
+    def get_agent_log(self, agent_id: AgentID) -> str:
+        if agent_id not in self._agent_logs:
+            raise UnknownRecordError("Error occured while getting agent")
+        return self._agent_logs[agent_id]
+
+    def reset(self):
+        self._agent_logs = {}

monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py

@@ -3,9 +3,9 @@ from uuid import UUID
 
 import pytest
 from tests.common import StubDIContainer
+from tests.monkey_island import InMemoryAgentLogRepository
 
-from common.types import AgentID
+from monkey_island.cc.repository import IAgentLogRepository
-from monkey_island.cc.repository import IAgentLogRepository, UnknownRecordError
 
 AGENT_ID_1 = UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5")
 AGENT_ID_2 = UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5")
@@ -14,27 +14,10 @@ AGENT_LOGS_URL_1 = f"/api/agent-logs/{AGENT_ID_1}"
 AGENT_LOGS_URL_2 = f"/api/agent-logs/{AGENT_ID_2}"
 
 
-class StubAgentLogRepository(IAgentLogRepository):
-    def __init__(self):
-        self._agent_logs = {}
-
-    def upsert_agent_log(self, agent_id: AgentID, log_contents: str):
-        if agent_id not in self._agent_logs.keys():
-            self._agent_logs[agent_id] = log_contents
-
-    def get_agent_log(self, agent_id: AgentID) -> str:
-        if agent_id not in self._agent_logs:
-            raise UnknownRecordError("Error occured while getting agent")
-        return self._agent_logs[agent_id]
-
-    def reset(self):
-        self._agent_logs = {}
-
-
 @pytest.fixture
 def flask_client(build_flask_client):
     container = StubDIContainer()
-    container.register_instance(IAgentLogRepository, StubAgentLogRepository())
+    container.register_instance(IAgentLogRepository, InMemoryAgentLogRepository())
 
     with build_flask_client(container) as flask_client:
         yield flask_client