diff --git a/monkey/tests/unit_tests/infection_monkey/exploit/test_zerologon.py b/monkey/tests/unit_tests/infection_monkey/exploit/test_zerologon.py index d3854d57c..6f5f9bbd6 100644 --- a/monkey/tests/unit_tests/infection_monkey/exploit/test_zerologon.py +++ b/monkey/tests/unit_tests/infection_monkey/exploit/test_zerologon.py @@ -16,7 +16,12 @@ NT_HASHES = ["def456", "765vut"] @pytest.fixture -def zerologon_exploiter_object(monkeypatch): +def mock_agent_event_queue(): + return MagicMock(spec=IAgentEventQueue) + + +@pytest.fixture +def zerologon_exploiter_object(monkeypatch, mock_agent_event_queue): from infection_monkey.exploit.zerologon import ZerologonExploiter def mock_report_login_attempt(**kwargs): @@ -26,22 +31,26 @@ def zerologon_exploiter_object(monkeypatch): monkeypatch.setattr(obj, "dc_name", NETBIOS_NAME, raising=False) monkeypatch.setattr(obj, "report_login_attempt", mock_report_login_attempt) monkeypatch.setattr(obj, "host", VictimHost(ip_addr="1.1.1.1")) - monkeypatch.setattr(obj, "agent_event_queue", MagicMock(spec=IAgentEventQueue)) + monkeypatch.setattr(obj, "agent_event_queue", mock_agent_event_queue) return obj @pytest.mark.slow -def test_assess_exploit_attempt_result_no_error(zerologon_exploiter_object): +def test_assess_exploit_attempt_result_no_error(zerologon_exploiter_object, mock_agent_event_queue): dummy_exploit_attempt_result = {"ErrorCode": 0} assert zerologon_exploiter_object.assess_exploit_attempt_result(dummy_exploit_attempt_result) + assert mock_agent_event_queue.publish.call_count == 1 @pytest.mark.slow -def test_assess_exploit_attempt_result_with_error(zerologon_exploiter_object): +def test_assess_exploit_attempt_result_with_error( + zerologon_exploiter_object, mock_agent_event_queue +): dummy_exploit_attempt_result = {"ErrorCode": 1} assert not zerologon_exploiter_object.assess_exploit_attempt_result( dummy_exploit_attempt_result ) + assert mock_agent_event_queue.publish.call_count == 1 @pytest.mark.slow