From 1a2a4d165a5905187cfebe7a7e1513fef8725b55 Mon Sep 17 00:00:00 2001
From: Shreya Malviya <shreya.malviya@gmail.com>
Date: Mon, 8 Aug 2022 23:56:10 +0530
Subject: [PATCH] Common: Add tests for PypubsubEventQueue

All the unsubscribing tests fails for some reason right now. The
subscriber is unsubscribed from the specified topics but the subscriber
is still called when an event with that topic is published.
---
 .../event_queue/test_pypubsub_event_queue.py  | 105 ++++++++++++++++++
 1 file changed, 105 insertions(+)
 create mode 100644 monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py

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
new file mode 100644
index 000000000..b06c37c86
--- /dev/null
+++ b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_event_queue.py
@@ -0,0 +1,105 @@
+from dataclasses import dataclass
+from unittest.mock import MagicMock
+
+import pytest
+
+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
+
+EVENT_TAG_1 = "event tag 1"
+EVENT_TAG_2 = "event tag 2"
+
+
+@dataclass(frozen=True)
+class EventType(AbstractEvent):
+    source = "1234"
+    target = None
+    timestamp = 0.0
+    tags = [EVENT_TAG_1, EVENT_TAG_2]
+
+
+@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
+    )
+
+
+def test_subscribe_all():
+    subscriber = MagicMock()
+
+    PypubsubEventQueue.subscribe_all(subscriber)
+    PypubsubEventQueue.publish(EventType)
+
+    assert pypubsub_event_queue_file.pub.sendMessage.call_count == 3
+    assert subscriber.call_count == 3
+
+
+def test_subscribe_types():
+    subscriber = MagicMock()
+
+    PypubsubEventQueue.subscribe_types([EventType], subscriber)
+    PypubsubEventQueue.publish(EventType)
+
+    assert pypubsub_event_queue_file.pub.sendMessage.call_count == 3
+    assert subscriber.call_count == 1
+
+
+def test_subscribe_tags():
+    subscriber = MagicMock()
+
+    PypubsubEventQueue.subscribe_tags([EVENT_TAG_2], subscriber)
+    PypubsubEventQueue.publish(EventType)
+
+    assert pypubsub_event_queue_file.pub.sendMessage.call_count == 3
+    assert subscriber.call_count == 1
+
+
+def test_unsubscribe_all():
+    subscriber = MagicMock()
+
+    PypubsubEventQueue.subscribe_all(subscriber)
+    PypubsubEventQueue.publish(EventType)
+
+    assert pypubsub_event_queue_file.pub.sendMessage.call_count == 3
+    assert subscriber.call_count == 3
+
+    PypubsubEventQueue.unsubscribe_all(subscriber)
+    PypubsubEventQueue.publish(EventType)
+
+    assert pypubsub_event_queue_file.pub.sendMessage.call_count == 6
+    assert subscriber.call_count == 3
+
+
+def test_unsubscribe_types():
+    subscriber = MagicMock()
+
+    PypubsubEventQueue.subscribe_all(subscriber)
+    PypubsubEventQueue.publish(EventType)
+
+    assert pypubsub_event_queue_file.pub.sendMessage.call_count == 3
+    assert subscriber.call_count == 3
+
+    PypubsubEventQueue.unsubscribe_types([EventType], subscriber)
+    PypubsubEventQueue.publish(EventType)
+
+    assert pypubsub_event_queue_file.pub.sendMessage.call_count == 6
+    assert subscriber.call_count == 5
+
+
+def test_unsubscribe_tags():
+    subscriber = MagicMock()
+
+    PypubsubEventQueue.subscribe_all(subscriber)
+    PypubsubEventQueue.publish(EventType)
+
+    assert pypubsub_event_queue_file.pub.sendMessage.call_count == 3
+    assert subscriber.call_count == 3
+
+    PypubsubEventQueue.unsubscribe_tags([EVENT_TAG_1, EVENT_TAG_2], subscriber)
+    PypubsubEventQueue.publish(EventType)
+
+    assert pypubsub_event_queue_file.pub.sendMessage.call_count == 6
+    assert subscriber.call_count == 4