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 a7c6da202..6bae3746e 100644 --- a/monkey/tests/unit_tests/infection_monkey/exploit/test_zerologon.py +++ b/monkey/tests/unit_tests/infection_monkey/exploit/test_zerologon.py @@ -2,6 +2,7 @@ from unittest.mock import MagicMock import pytest +from common.agent_events import ExploitationEvent, PasswordRestorationEvent from common.event_queue import IAgentEventQueue from infection_monkey.model import VictimHost @@ -38,6 +39,7 @@ def test_assess_exploit_attempt_result_no_error(zerologon_exploiter_object, mock 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 + assert mock_agent_event_queue.publish.call_args[0][0].__class__ is ExploitationEvent assert mock_agent_event_queue.publish.call_args[0][0].success @@ -50,23 +52,34 @@ def test_assess_exploit_attempt_result_with_error( dummy_exploit_attempt_result ) assert mock_agent_event_queue.publish.call_count == 1 + assert mock_agent_event_queue.publish.call_args[0][0].__class__ is ExploitationEvent assert not mock_agent_event_queue.publish.call_args[0][0].success @pytest.mark.slow -def test_assess_restoration_attempt_result_restored(zerologon_exploiter_object): +def test_assess_restoration_attempt_result_restored( + zerologon_exploiter_object, mock_agent_event_queue +): dummy_restoration_attempt_result = object() assert zerologon_exploiter_object.assess_restoration_attempt_result( dummy_restoration_attempt_result ) + assert mock_agent_event_queue.publish.call_count == 1 + assert mock_agent_event_queue.publish.call_args[0][0].__class__ is PasswordRestorationEvent + assert mock_agent_event_queue.publish.call_args[0][0].success @pytest.mark.slow -def test_assess_restoration_attempt_result_not_restored(zerologon_exploiter_object): +def test_assess_restoration_attempt_result_not_restored( + zerologon_exploiter_object, mock_agent_event_queue +): dummy_restoration_attempt_result = False assert not zerologon_exploiter_object.assess_restoration_attempt_result( dummy_restoration_attempt_result ) + assert mock_agent_event_queue.publish.call_count == 1 + assert mock_agent_event_queue.publish.call_args[0][0].__class__ is PasswordRestorationEvent + assert not mock_agent_event_queue.publish.call_args[0][0].success @pytest.mark.slow