From 2d42355e2c990cc602e945b0650f7bbdce52b535 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Fri, 23 Sep 2022 12:21:36 +0200
Subject: [PATCH] UT: Add tests for GET AgentSignals endpoint

---
 .../cc/resources/test_agent_signals.py        | 50 ++++++++++++-------
 1 file changed, 33 insertions(+), 17 deletions(-)

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