From 3da90223fcabee9e288db5d3f83aa42a5a7cd557 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski 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()