UT: Improve tests for zerologon

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