From e60297dff1012702d060686ce6457cf8fdbe160f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 16 Dec 2021 14:36:33 -0500 Subject: [PATCH] UT: Fix broken mock in test_stop_if_cant_get_stop_signal_from_island --- .../master/test_automated_master.py | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py b/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py index 378eab883..0916acda7 100644 --- a/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py +++ b/monkey/tests/unit_tests/infection_monkey/master/test_automated_master.py @@ -1,6 +1,8 @@ import time from unittest.mock import MagicMock +import pytest + from infection_monkey.master import AutomatedMaster from infection_monkey.master.automated_master import ( CHECK_FOR_CONFIG_COUNT, @@ -38,19 +40,29 @@ def test_stop_if_cant_get_config_from_island(monkeypatch): assert cc.get_config.call_count == CHECK_FOR_CONFIG_COUNT +@pytest.fixture +def sleep_and_return_config(automated_master_config): + # Ensure that should_agent_stop times out before get_config() returns to prevent the + # Propagator's sub-threads from hanging + get_config_sleep_time = INTERVAL * (CHECK_FOR_STOP_AGENT_COUNT + 1) + + def _inner(): + time.sleep(get_config_sleep_time) + return automated_master_config + + return _inner + + # NOTE: This test is a little bit brittle, and probably needs too much knowlegde of the internals # of AutomatedMaster. For now, it works and it runs quickly. In the future, if we find that # this test isn't valuable or it starts causing issues, we can just remove it. -def test_stop_if_cant_get_stop_signal_from_island(monkeypatch, automated_master_config): +def test_stop_if_cant_get_stop_signal_from_island(monkeypatch, sleep_and_return_config): cc = MagicMock() cc.should_agent_stop = MagicMock( side_effect=IslandCommunicationError("Failed to communicate with island") ) - # Ensure that should_agent_stop times out before get_config() returns to prevent the - # Propagator's sub-threads from hanging cc.get_config = MagicMock( - return_value=automated_master_config, - side_effect=lambda: time.sleep(INTERVAL * (CHECK_FOR_STOP_AGENT_COUNT + 1)), + side_effect=sleep_and_return_config, ) monkeypatch.setattr(