From 3da90223fcabee9e288db5d3f83aa42a5a7cd557 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Fri, 23 Sep 2022 13:52:26 +0200
Subject: [PATCH] Agent: Use IIslandAPIClient.get_agent_signals in
ControlChannel
---
monkey/infection_monkey/master/control_channel.py | 2 +-
.../infection_monkey/master/test_control_channel.py | 4 ++--
2 files changed, 3 insertions(+), 3 deletions(-)
diff --git a/monkey/infection_monkey/master/control_channel.py b/monkey/infection_monkey/master/control_channel.py
index 48b827f63..947d6c0da 100644
--- a/monkey/infection_monkey/master/control_channel.py
+++ b/monkey/infection_monkey/master/control_channel.py
@@ -36,7 +36,7 @@ class ControlChannel(IControlChannel):
if not self._control_channel_server:
logger.error("Agent should stop because it can't connect to the C&C server.")
return True
- return self._island_api_client.should_agent_stop(self._agent_id)
+ return self._island_api_client.get_agent_signals(self._agent_id) is not None
@handle_island_api_errors
def get_config(self) -> AgentConfiguration:
diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py
index 658635615..1da0d0713 100644
--- a/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py
+++ b/monkey/tests/unit_tests/infection_monkey/master/test_control_channel.py
@@ -35,14 +35,14 @@ def control_channel(island_api_client) -> ControlChannel:
def test_control_channel__should_agent_stop(control_channel, island_api_client):
control_channel.should_agent_stop()
- assert island_api_client.should_agent_stop.called_once()
+ assert island_api_client.get_agent_signals.called_once()
@pytest.mark.parametrize("api_error", CONTROL_CHANNEL_API_ERRORS)
def test_control_channel__should_agent_stop_raises_error(
control_channel, island_api_client, api_error
):
- island_api_client.should_agent_stop.side_effect = api_error()
+ island_api_client.get_agent_signals.side_effect = api_error()
with pytest.raises(IslandCommunicationError):
control_channel.should_agent_stop()