From 28cf860aeb581196a1686507c5a88f7822957428 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 5 Sep 2022 19:13:58 +0530 Subject: [PATCH] UT: Use IAgentEventQueue in place of IEventQueue --- .../event_queue/test_pypubsub_event_queue.py | 16 ++++++++-------- .../test_mimikatz_collector.py | 10 +++++----- .../test_ssh_credentials_collector.py | 6 +++--- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py index 0da3c9f01..8e5bfe64e 100644 --- a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py +++ b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py @@ -6,7 +6,7 @@ from uuid import UUID import pytest from pubsub.core import Publisher -from common.event_queue import EventSubscriber, IEventQueue, PyPubSubEventQueue +from common.event_queue import EventSubscriber, IAgentEventQueue, PyPubSubEventQueue from common.events import AbstractAgentEvent EVENT_TAG_1 = "event tag 1" @@ -32,7 +32,7 @@ class TestEvent2(AbstractAgentEvent): @pytest.fixture -def event_queue() -> IEventQueue: +def event_queue() -> IAgentEventQueue: return PyPubSubEventQueue(Publisher()) @@ -50,7 +50,7 @@ def event_queue_subscriber() -> Callable[[AbstractAgentEvent], None]: return fn -def test_subscribe_all(event_queue: IEventQueue, event_queue_subscriber: EventSubscriber): +def test_subscribe_all(event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber): event_queue.subscribe_all_events(event_queue_subscriber) event_queue.publish(TestEvent1(tags=frozenset({EVENT_TAG_1, EVENT_TAG_2}))) @@ -65,7 +65,7 @@ def test_subscribe_all(event_queue: IEventQueue, event_queue_subscriber: EventSu @pytest.mark.parametrize("type_to_subscribe", [TestEvent1, TestEvent2]) def test_subscribe_types( - event_queue: IEventQueue, event_queue_subscriber: EventSubscriber, type_to_subscribe + event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber, type_to_subscribe ): event_queue.subscribe_type(type_to_subscribe, event_queue_subscriber) @@ -77,7 +77,7 @@ def test_subscribe_types( def test_subscribe_tags_single_type( - event_queue: IEventQueue, event_queue_subscriber: EventSubscriber + event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber ): event_queue.subscribe_tag(EVENT_TAG_1, event_queue_subscriber) @@ -91,7 +91,7 @@ def test_subscribe_tags_single_type( def test_subscribe_tags_multiple_types( - event_queue: IEventQueue, event_queue_subscriber: EventSubscriber + event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber ): event_queue.subscribe_tag(EVENT_TAG_2, event_queue_subscriber) @@ -105,7 +105,7 @@ def test_subscribe_tags_multiple_types( assert {EVENT_TAG_1, EVENT_TAG_2}.issubset(event_queue_subscriber.call_tags) -def test_type_tag_collision(event_queue: IEventQueue, event_queue_subscriber: EventSubscriber): +def test_type_tag_collision(event_queue: IAgentEventQueue, event_queue_subscriber: EventSubscriber): event_queue.subscribe_type(TestEvent1, event_queue_subscriber) event_queue.publish(TestEvent2(tags=frozenset({TestEvent1.__name__}))) @@ -113,7 +113,7 @@ def test_type_tag_collision(event_queue: IEventQueue, event_queue_subscriber: Ev assert event_queue_subscriber.call_count == 0 -def test_keep_subscriber_in_scope(event_queue: IEventQueue): +def test_keep_subscriber_in_scope(event_queue: IAgentEventQueue): class MyCallable: called = False 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 12ca9d594..0476d8a13 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 @@ -4,7 +4,7 @@ from unittest.mock import MagicMock import pytest from common.credentials import Credentials, LMHash, NTHash, Password, Username -from common.event_queue import IEventQueue +from common.event_queue import IAgentEventQueue 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 @@ -24,7 +24,7 @@ def patch_pypykatz(win_creds: [WindowsCredentials], monkeypatch): def collect_credentials() -> Sequence[Credentials]: - mock_event_queue = MagicMock(spec=IEventQueue) + mock_event_queue = MagicMock(spec=IAgentEventQueue) return MimikatzCredentialCollector(mock_event_queue).collect_credentials() @@ -120,7 +120,7 @@ def test_pypykatz_result_parsing_no_secrets(monkeypatch): def test_mimikatz_credentials_stolen_event_published(monkeypatch): - mock_event_queue = MagicMock(spec=IEventQueue) + mock_event_queue = MagicMock(spec=IAgentEventQueue) patch_pypykatz([], monkeypatch) mimikatz_credential_collector = MimikatzCredentialCollector(mock_event_queue) @@ -134,7 +134,7 @@ def test_mimikatz_credentials_stolen_event_published(monkeypatch): def test_mimikatz_credentials_stolen_event_tags(monkeypatch): - mock_event_queue = MagicMock(spec=IEventQueue) + mock_event_queue = MagicMock(spec=IAgentEventQueue) patch_pypykatz([], monkeypatch) mimikatz_credential_collector = MimikatzCredentialCollector(mock_event_queue) @@ -146,7 +146,7 @@ def test_mimikatz_credentials_stolen_event_tags(monkeypatch): def test_mimikatz_credentials_stolen_event_stolen_credentials(monkeypatch): - mock_event_queue = MagicMock(spec=IEventQueue) + mock_event_queue = MagicMock(spec=IAgentEventQueue) win_creds = [ WindowsCredentials( username="user2", password="secret2", lm_hash="0182BD0BD4444BF8FC83B5D9042EED2E" diff --git a/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_ssh_credentials_collector.py b/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_ssh_credentials_collector.py index ba12e416f..1c7fbe1b4 100644 --- a/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_ssh_credentials_collector.py +++ b/monkey/tests/unit_tests/infection_monkey/credential_collectors/test_ssh_credentials_collector.py @@ -3,7 +3,7 @@ from unittest.mock import MagicMock import pytest from common.credentials import Credentials, SSHKeypair, Username -from common.event_queue import IEventQueue +from common.event_queue import IAgentEventQueue from infection_monkey.credential_collectors import SSHCredentialCollector @@ -25,7 +25,7 @@ def patch_ssh_handler(ssh_creds, monkeypatch): def test_ssh_credentials_empty_results(monkeypatch, ssh_creds, patch_telemetry_messenger): patch_ssh_handler(ssh_creds, monkeypatch) collected = SSHCredentialCollector( - patch_telemetry_messenger, MagicMock(spec=IEventQueue) + patch_telemetry_messenger, MagicMock(spec=IAgentEventQueue) ).collect_credentials() assert not collected @@ -71,6 +71,6 @@ def test_ssh_info_result_parsing(monkeypatch, patch_telemetry_messenger): Credentials(identity=None, secret=ssh_keypair3), ] collected = SSHCredentialCollector( - patch_telemetry_messenger, MagicMock(spec=IEventQueue) + patch_telemetry_messenger, MagicMock(spec=IAgentEventQueue) ).collect_credentials() assert expected == collected