UT: Use IAgentEventQueue in place of IEventQueue

This commit is contained in:
Shreya Malviya 2022-09-05 19:13:58 +05:30
parent acb049f1f1
commit 28cf860aeb
3 changed files with 16 additions and 16 deletions

View File

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

View File

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

View File

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