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 738b91566..50c9635d9 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 @@ -2,8 +2,8 @@ from dataclasses import dataclass from unittest.mock import MagicMock import pytest +from pubsub import pub -import common.event_queue.pypubsub_event_queue as pypubsub_event_queue_file from common.event_queue.pypubsub_event_queue import PyPubSubEventQueue from common.events import AbstractEvent @@ -22,36 +22,38 @@ class EventType(AbstractEvent): @pytest.fixture(autouse=True) def wrap_pypubsub_functions(): # This is done so that we can use `.call_count` in the tests. - pypubsub_event_queue_file.pub.sendMessage = MagicMock( - side_effect=pypubsub_event_queue_file.pub.sendMessage - ) + pub.sendMessage = MagicMock(side_effect=pub.sendMessage) + + +pypubsub_event_queue = PyPubSubEventQueue(pub) def test_subscribe_all(): subscriber = MagicMock() - PyPubSubEventQueue.subscribe_all(subscriber) - PyPubSubEventQueue.publish(EventType) + pypubsub_event_queue.subscribe_all(subscriber) + pypubsub_event_queue.publish(EventType) - assert pypubsub_event_queue_file.pub.sendMessage.call_count == 3 + assert pub.sendMessage.call_count == 3 assert subscriber.call_count == 3 def test_subscribe_types(): subscriber = MagicMock() - PyPubSubEventQueue.subscribe_type(EventType, subscriber) - PyPubSubEventQueue.publish(EventType) + pypubsub_event_queue.subscribe_type(EventType, subscriber) + pypubsub_event_queue.publish(EventType) - assert pypubsub_event_queue_file.pub.sendMessage.call_count == 3 + assert pub.sendMessage.call_count == 3 assert subscriber.call_count == 1 def test_subscribe_tags(): subscriber = MagicMock() - PyPubSubEventQueue.subscribe_tag(EVENT_TAG_2, subscriber) - PyPubSubEventQueue.publish(EventType) + pypubsub_event_queue.subscribe_tag(EVENT_TAG_1, subscriber) + pypubsub_event_queue.subscribe_tag(EVENT_TAG_2, subscriber) + pypubsub_event_queue.publish(EventType) - assert pypubsub_event_queue_file.pub.sendMessage.call_count == 3 - assert subscriber.call_count == 1 + assert pub.sendMessage.call_count == 3 + assert subscriber.call_count == 2