UT: Add mock_agent_event_queue in zerologon tests

This commit is contained in:
Ilija Lazoroski 2022-10-05 17:04:30 +02:00
parent 9a261296e4
commit 1db5e9adc5
1 changed files with 13 additions and 4 deletions

View File

@ -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