diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py index 63ecef8f2..6fa371b09 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agent_signals.py @@ -6,29 +6,28 @@ import pytest from tests.common import StubDIContainer from monkey_island.cc.event_queue import IIslandEventQueue +from monkey_island.cc.models import AgentSignals as Signals +from monkey_island.cc.repository import RetrievalError, StorageError from monkey_island.cc.resources import AgentSignals from monkey_island.cc.services import AgentSignalsService TIMESTAMP = 123456789 +TIMESTAMP_1 = 123546789 +SIGNALS = Signals(terminate=TIMESTAMP) +SIGNALS_1 = Signals(terminate=TIMESTAMP_1) -@pytest.fixture( - params=[ - UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5"), - UUID("9b4279f6-6ec5-4953-821e-893ddc71a988"), - ] -) -def agent_id(request) -> UUID: - return request.param +AGENT_ID = UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5") +AGENT_ID_1 = UUID("9b4279f6-6ec5-4953-821e-893ddc71a988") @pytest.fixture -def agent_signals_url(agent_id: UUID) -> str: - return f"/api/agent-signals/{agent_id}" +def mock_agent_signals_service(): + return MagicMock(spec=AgentSignalsService) @pytest.fixture -def flask_client_builder(build_flask_client): +def flask_client_builder(build_flask_client, mock_agent_signals_service): def inner(side_effect=None): container = StubDIContainer() @@ -36,7 +35,6 @@ def flask_client_builder(build_flask_client): mock_island_event_queue.publish.side_effect = side_effect container.register_instance(IIslandEventQueue, mock_island_event_queue) - mock_agent_signals_service = MagicMock(spec=AgentSignalsService) container.register_instance(AgentSignalsService, mock_agent_signals_service) with build_flask_client(container) as flask_client: @@ -77,9 +75,27 @@ def test_agent_signals_terminate_all_post__invalid_timestamp(flask_client, bad_d assert resp.status_code == HTTPStatus.BAD_REQUEST -# TODO: Complete this when GET is implemented -# Do we get a value indicating that we should stop? Depends on whether a signal was sent -def test_agent_signals_endpoint(flask_client, agent_signals_url): - resp = flask_client.get(agent_signals_url, follow_redirects=True) +@pytest.mark.parametrize( + "url, signals", + [(f"/api/agent-signals/{AGENT_ID}", SIGNALS), (f"/api/agent-signals/{AGENT_ID_1}", SIGNALS_1)], +) +def test_agent_signals_get(flask_client, mock_agent_signals_service, url, signals): + mock_agent_signals_service.get_signals.return_value = signals + resp = flask_client.get(url, follow_redirects=True) assert resp.status_code == HTTPStatus.OK - assert resp.json == {} + assert resp.json == signals.dict(simplify=True) + + +@pytest.mark.parametrize( + "url, error", + [ + (f"/api/agent-signals/{AGENT_ID}", RetrievalError), + (f"/api/agent-signals/{AGENT_ID_1}", StorageError), + ], +) +def test_agent_signals_get__internal_server_error( + flask_client, mock_agent_signals_service, url, error +): + mock_agent_signals_service.get_signals.side_effect = error + resp = flask_client.get(url, follow_redirects=True) + assert resp.status_code == HTTPStatus.INTERNAL_SERVER_ERROR