UT: Use event_queue_subscriber fixture in MimikatzCredentialCollector tests

This commit is contained in:
Shreya Malviya 2022-08-16 19:16:25 +05:30
parent b6c7001294
commit 140967b05f
1 changed files with 6 additions and 12 deletions

View File

@ -6,7 +6,6 @@ from pubsub.core import Publisher
from common.credentials import Credentials, LMHash, NTHash, Password, Username
from common.event_queue import IEventQueue, PyPubSubEventQueue
from common.events import AbstractEvent
from infection_monkey.credential_collectors import MimikatzCredentialCollector
from infection_monkey.credential_collectors.mimikatz_collector.mimikatz_credential_collector import ( # noqa: E501
MIMIKATZ_EVENT_TAGS,
@ -125,16 +124,11 @@ def event_queue() -> IEventQueue:
return PyPubSubEventQueue(Publisher())
def test_pypykatz_credentials_stolen_event_published(monkeypatch, event_queue):
def subscriber(event: AbstractEvent):
subscriber.call_count += 1
subscriber.call_tags |= event.tags
subscriber.call_count = 0
subscriber.call_tags = set()
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, subscriber)
event_queue.subscribe_tag(event_tag, event_queue_subscriber)
mimikatz_credential_collector = MimikatzCredentialCollector(event_queue)
monkeypatch.setattr(
@ -142,5 +136,5 @@ def test_pypykatz_credentials_stolen_event_published(monkeypatch, event_queue):
)
mimikatz_credential_collector.collect_credentials()
assert subscriber.call_count == 3
assert subscriber.call_tags == MIMIKATZ_EVENT_TAGS
assert event_queue_subscriber.call_count == 3
assert event_queue_subscriber.call_tags == MIMIKATZ_EVENT_TAGS