UT: Add tests for GET AgentSignals endpoint

This commit is contained in:
Ilija Lazoroski 2022-09-23 12:21:36 +02:00 committed by Mike Salvatore
parent 41951511d0
commit 2d42355e2c
1 changed files with 33 additions and 17 deletions

View File

@ -6,29 +6,28 @@ import pytest
from tests.common import StubDIContainer from tests.common import StubDIContainer
from monkey_island.cc.event_queue import IIslandEventQueue 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.resources import AgentSignals
from monkey_island.cc.services import AgentSignalsService from monkey_island.cc.services import AgentSignalsService
TIMESTAMP = 123456789 TIMESTAMP = 123456789
TIMESTAMP_1 = 123546789
SIGNALS = Signals(terminate=TIMESTAMP)
SIGNALS_1 = Signals(terminate=TIMESTAMP_1)
@pytest.fixture( AGENT_ID = UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5")
params=[ AGENT_ID_1 = UUID("9b4279f6-6ec5-4953-821e-893ddc71a988")
UUID("c0dd10b3-e21a-4da9-9d96-a99c19ebd7c5"),
UUID("9b4279f6-6ec5-4953-821e-893ddc71a988"),
]
)
def agent_id(request) -> UUID:
return request.param
@pytest.fixture @pytest.fixture
def agent_signals_url(agent_id: UUID) -> str: def mock_agent_signals_service():
return f"/api/agent-signals/{agent_id}" return MagicMock(spec=AgentSignalsService)
@pytest.fixture @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): def inner(side_effect=None):
container = StubDIContainer() container = StubDIContainer()
@ -36,7 +35,6 @@ def flask_client_builder(build_flask_client):
mock_island_event_queue.publish.side_effect = side_effect mock_island_event_queue.publish.side_effect = side_effect
container.register_instance(IIslandEventQueue, mock_island_event_queue) container.register_instance(IIslandEventQueue, mock_island_event_queue)
mock_agent_signals_service = MagicMock(spec=AgentSignalsService)
container.register_instance(AgentSignalsService, mock_agent_signals_service) container.register_instance(AgentSignalsService, mock_agent_signals_service)
with build_flask_client(container) as flask_client: 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 assert resp.status_code == HTTPStatus.BAD_REQUEST
# TODO: Complete this when GET is implemented @pytest.mark.parametrize(
# Do we get a value indicating that we should stop? Depends on whether a signal was sent "url, signals",
def test_agent_signals_endpoint(flask_client, agent_signals_url): [(f"/api/agent-signals/{AGENT_ID}", SIGNALS), (f"/api/agent-signals/{AGENT_ID_1}", SIGNALS_1)],
resp = flask_client.get(agent_signals_url, follow_redirects=True) )
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.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