From 6299529f4a5dbad50afd9645cb483daa2830b8bb Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 17:11:44 +0200 Subject: [PATCH] Island: Modify HTTPIslandAPIClient.get_agent_signals to return AgentSignals --- .../island_api_client/http_island_api_client.py | 9 ++++----- .../test_http_island_api_client.py | 13 +++++++------ 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index 65e34df3c..09fcbf762 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -1,13 +1,12 @@ import functools import json import logging -from datetime import datetime from pprint import pformat -from typing import List, Optional, Sequence +from typing import List, Sequence import requests -from common import AgentRegistrationData, OperatingSystem +from common import AgentRegistrationData, AgentSignals, OperatingSystem from common.agent_configuration import AgentConfiguration from common.agent_event_serializers import AgentEventSerializerRegistry, JSONSerializable from common.agent_events import AbstractAgentEvent @@ -189,7 +188,7 @@ class HTTPIslandAPIClient(IIslandAPIClient): @handle_island_errors @convert_json_error_to_island_api_error - def get_agent_signals(self, agent_id: str) -> Optional[datetime]: + def get_agent_signals(self, agent_id: str) -> AgentSignals: url = f"{self._api_url}/agent-signals/{agent_id}" response = requests.get( # noqa: DUO123 url, @@ -197,7 +196,7 @@ class HTTPIslandAPIClient(IIslandAPIClient): timeout=SHORT_REQUEST_TIMEOUT, ) response.raise_for_status() - return response.json()["terminate"] + return AgentSignals(**response.json()) class HTTPIslandAPIClientFactory(AbstractIslandAPIClientFactory): diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index 9505e6649..a861a4849 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -4,7 +4,7 @@ import pytest import requests import requests_mock -from common import OperatingSystem +from common import AgentSignals, OperatingSystem from common.agent_event_serializers import ( AgentEventSerializerRegistry, PydanticAgentEventSerializer, @@ -456,16 +456,17 @@ def test_island_api_client_get_agent_signals__status_code( island_api_client.get_agent_signals(agent_id=AGENT_ID) -@pytest.mark.parametrize("expected_timestamp", [TIMESTAMP, None]) -def test_island_api_client_get_agent_signals(island_api_client, expected_timestamp): +@pytest.mark.parametrize("timestamp", [TIMESTAMP, None]) +def test_island_api_client_get_agent_signals(island_api_client, timestamp): + expected_agent_signals = AgentSignals(terminate=timestamp) with requests_mock.Mocker() as m: m.get(ISLAND_URI) island_api_client.connect(SERVER) - m.get(ISLAND_GET_AGENT_SIGNALS, json={"terminate": expected_timestamp}) - actual_terminate_timestamp = island_api_client.get_agent_signals(agent_id=AGENT_ID) + m.get(ISLAND_GET_AGENT_SIGNALS, json={"terminate": timestamp}) + actual_agent_signals = island_api_client.get_agent_signals(agent_id=AGENT_ID) - assert actual_terminate_timestamp == expected_timestamp + assert actual_agent_signals == expected_agent_signals def test_island_api_client_get_agent_signals__bad_json(island_api_client):