diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py new file mode 100644 index 000000000..ead1c73a2 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_logs.py @@ -0,0 +1,59 @@ +from http import HTTPStatus +from uuid import UUID + +import pytest +from tests.common import StubDIContainer + +from common.types import AgentID +from monkey_island.cc.repository import IAgentLogRepository + +AGENT_ID_1 = UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5") +AGENT_ID_2 = UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5") + +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: + 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()) + + with build_flask_client(container) as flask_client: + yield flask_client + + +def test_agent_logs_endpoint__get_empty(flask_client): + resp = flask_client.get(AGENT_LOGS_URL_1, follow_redirects=True) + assert resp.status_code == HTTPStatus.INTERNAL_SERVER_ERROR + + +@pytest.mark.parametrize( + "url,log", [(AGENT_LOGS_URL_1, "LoremIpsum1"), (AGENT_LOGS_URL_2, "SecondLoremIpsum")] +) +def test_agent_logs_endpoint(flask_client, url, log): + flask_client.put(url, json={"log_contents": log}, follow_redirects=True) + resp = flask_client.get(url, follow_redirects=True) + assert resp.status_code == HTTPStatus.OK + assert resp.json == log + + +def test_agent_logs_endpoint__put_bad_json(flask_client): + resp = flask_client.put(AGENT_LOGS_URL_1, json={"bad_key": "bad"}, follow_redirects=True) + + assert resp.status_code == HTTPStatus.INTERNAL_SERVER_ERROR