From eb17b2062503fb69159b59e137f5bbe88efde3bd Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 17 Aug 2022 14:41:26 +0530 Subject: [PATCH] UT: Simplify MimikatzCredentialCollector's event publishing test --- .../test_mimikatz_collector.py | 35 ++++++++++--------- 1 file changed, 18 insertions(+), 17 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_mimikatz_collector.py b/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_mimikatz_collector.py index 0910aab1b..7ad4a1834 100644 --- a/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_mimikatz_collector.py +++ b/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_mimikatz_collector.py @@ -2,10 +2,10 @@ from typing import Sequence from unittest.mock import MagicMock, Mock import pytest -from pubsub.core import Publisher from common.credentials import Credentials, LMHash, NTHash, Password, Username -from common.event_queue import IEventQueue, PyPubSubEventQueue +from common.event_queue import IEventQueue +from common.events import CredentialsStolenEvent from infection_monkey.credential_collectors import MimikatzCredentialCollector from infection_monkey.credential_collectors.mimikatz_collector.mimikatz_credential_collector import ( # noqa: E501 MIMIKATZ_EVENT_TAGS, @@ -119,22 +119,23 @@ def test_pypykatz_result_parsing_no_secrets(monkeypatch): assert collected_credentials == expected_credentials -@pytest.fixture -def event_queue() -> IEventQueue: - return PyPubSubEventQueue(Publisher()) - - -def test_pypykatz_credentials_stolen_event_published( - monkeypatch, event_queue, event_queue_subscriber -): - for event_tag in MIMIKATZ_EVENT_TAGS: - event_queue.subscribe_tag(event_tag, event_queue_subscriber) - - mimikatz_credential_collector = MimikatzCredentialCollector(event_queue) +def test_mimikatz_credentials_stolen_event_published(monkeypatch): + mock_event_queue = MagicMock(spec=IEventQueue) + mock_subscriber = Mock() monkeypatch.setattr( "infection_monkey.credential_collectors.mimikatz_collector.pypykatz_handler", Mock() ) - mimikatz_credential_collector.collect_credentials() - assert event_queue_subscriber.call_count == 3 - assert event_queue_subscriber.call_tags == MIMIKATZ_EVENT_TAGS + for event_tag in MIMIKATZ_EVENT_TAGS: + mock_event_queue.subscribe_tag(event_tag, mock_subscriber) + + mimikatz_credential_collector = MimikatzCredentialCollector(mock_event_queue) + collected_credentials = mimikatz_credential_collector.collect_credentials() + + mock_event_queue.publish.assert_called_once() + + mock_event_queue_call_args = mock_event_queue.publish.call_args[0][0] + + assert isinstance(mock_event_queue_call_args, CredentialsStolenEvent) + assert mock_event_queue_call_args.tags == frozenset(MIMIKATZ_EVENT_TAGS) + assert mock_event_queue_call_args.stolen_credentials == collected_credentials