Island: Modify HTTPIslandAPIClient.get_agent_signals to return

AgentSignals

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):

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):