From d6501ebe0f1276c4b3ea62dcfab92bd7caea3076 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Tue, 12 Sep 2023 18:43:26 +0100
Subject: [PATCH] Add test for .peek()

---
 README.md                        |  2 +-
 permchain/connection_inmemory.py |  4 +--
 tests/test_invoke.py             | 44 ++++++++++++++++++++++++++++++--
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 4ceeb628c..c18df4c87 100644
--- a/README.md
+++ b/README.md
@@ -28,7 +28,7 @@ Check `tests` and `examples` for more examples.
 - [x] Implement IN as regular topic
 - [x] Add Connection.peek() to monitor past messages from all topics
 - [x] Enable resuming PubSub from the "middle" of the computation
-- [ ] Add test for .peek()
+- [x] Add test for .peek()
 - [ ] Move tracking of inflight processes/messages to Connection
   - [ ] Use this to build retry mechanism, where any inflight messages are moved back to the respective topics when restarting
 - [ ] Detect cycles (aka. infinite loops) and throw an error
diff --git a/permchain/connection_inmemory.py b/permchain/connection_inmemory.py
index 0762a3d89..0fdf7f86c 100644
--- a/permchain/connection_inmemory.py
+++ b/permchain/connection_inmemory.py
@@ -35,7 +35,7 @@ class InMemoryPubSubConnection(PubSubConnection):
         self.lock = threading.RLock()
 
     def peek(self, prefix: str) -> Iterator[LogMessage]:
-        return iter(self.logs[prefix])
+        return iter(self.logs[str(prefix)])
 
     def iterate(self, prefix: str, topic_name: str) -> Iterator[Any]:
         topic = self.full_topic_name(prefix, topic_name)
@@ -70,7 +70,7 @@ class InMemoryPubSubConnection(PubSubConnection):
 
         with self.lock:
             # Add the message to the log
-            self.logs[prefix].append(
+            self.logs[str(prefix)].append(
                 LogMessage(
                     message=message,
                     topic_name=topic_name,
diff --git a/tests/test_invoke.py b/tests/test_invoke.py
index d0f48c29d..3bad75f7a 100644
--- a/tests/test_invoke.py
+++ b/tests/test_invoke.py
@@ -68,10 +68,50 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture):
     # Then invoke both pubsubs, as a group
     # The second picks up where the first left off
     with trace_as_chain_group("PubSubGroup") as cm:
+        # invoke() step 1
         assert pubsub_one.invoke(2, {"callbacks": cm}) == []
+
+        # listeners are still cleared, even though state is preserved
+        assert conn.listeners == {}
+        # The log contains all messages published to all topics, in order
+        assert [{**m, "started_at": None} for m in conn.peek(cm.parent_run_id)] == [
+            {"message": 2, "topic_name": "__in__", "started_at": None},
+            {"message": 3, "topic_name": "one", "started_at": None},
+        ]
+        # IN, OUT, one
+        assert len(conn.topics) == 3
+        topic_one_full_name = conn.full_topic_name(cm.parent_run_id, topic_one.name)
+        # the actual message publishd by chain_one, and a sentinel "end" value
+        assert conn.topics[topic_one_full_name].qsize() == 2
+
+        # invoke() step 2
+        # this picks up where the first left off, and produces same result as
+        # `test_invoke_two_processes_in_out`
         assert pubsub_two.invoke(None, {"callbacks": cm}) == [4]
-    # After invoke returns the listeners were cleaned up
-    assert conn.listeners == {}
+
+        # listeners are still cleared, even though state is preserved
+        assert conn.listeners == {}
+        # The log contains all messages published to all topics, in order
+        assert [{**m, "started_at": None} for m in conn.peek(cm.parent_run_id)] == [
+            {"message": 2, "topic_name": "__in__", "started_at": None},
+            {"message": 3, "topic_name": "one", "started_at": None},
+            {"message": None, "topic_name": "__in__", "started_at": None},
+            {"message": 4, "topic_name": "__out__", "started_at": None},
+        ]
+        # IN, OUT, one
+        assert len(conn.topics) == 3
+
+        for topic_name, queue in conn.topics.items():
+            if topic_name.endswith("IN"):
+                # Contains two sentinel "end" values, and None
+                # passed in as input to chain_two, which doesn't subscribe to it
+                assert queue.qsize() == 3
+            if topic_name.endswith("OUT"):
+                # Empty because this was consumed by invoke()
+                assert queue.qsize() == 0
+            if topic_name.endswith("one"):
+                # Contains sentinel "end" value
+                assert queue.qsize() == 1
 
 
 def test_invoke_many_processes_in_out(mocker: MockerFixture):