diff --git a/permchain/connection.py b/permchain/connection.py index 6380ac670..bfaa71efc 100644 --- a/permchain/connection.py +++ b/permchain/connection.py @@ -6,9 +6,10 @@ PubSubListener = Callable[[Any], None] class LogMessage(TypedDict): - message: Any - topic_name: str - started_at: str + topic: str + value: Any + published_at: str + correlation_id: str class PubSubConnection(ABC): diff --git a/permchain/connection_inmemory.py b/permchain/connection_inmemory.py index d18d2577d..82c98fb33 100644 --- a/permchain/connection_inmemory.py +++ b/permchain/connection_inmemory.py @@ -78,9 +78,10 @@ class InMemoryPubSubConnection(PubSubConnection): # Add the message to the log self.logs[str(prefix)].append( LogMessage( - message=message, - topic_name=topic_name, - started_at=datetime.now().isoformat(), + value=message, + topic=topic_name, + correlation_id=str(prefix), + published_at=datetime.now().isoformat(), ) ) listeners = self.listeners[topic] diff --git a/tests/test_invoke.py b/tests/test_invoke.py index 8f36e385d..8b2d5f5ef 100644 --- a/tests/test_invoke.py +++ b/tests/test_invoke.py @@ -76,9 +76,19 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture): # 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(correlation_id)] == [ - {"message": 2, "topic_name": "__in__", "started_at": None}, - {"message": 3, "topic_name": "one", "started_at": None}, + assert [{**m, "published_at": None} for m in conn.peek(correlation_id)] == [ + { + "value": 2, + "topic": "__in__", + "correlation_id": str(correlation_id), + "published_at": None, + }, + { + "value": 3, + "topic": "one", + "correlation_id": str(correlation_id), + "published_at": None, + }, ] # IN, OUT, one assert len(conn.topics) == 3 @@ -94,11 +104,31 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture): # 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(correlation_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}, + assert [{**m, "published_at": None} for m in conn.peek(correlation_id)] == [ + { + "value": 2, + "topic": "__in__", + "correlation_id": str(correlation_id), + "published_at": None, + }, + { + "value": 3, + "topic": "one", + "correlation_id": str(correlation_id), + "published_at": None, + }, + { + "value": None, + "topic": "__in__", + "correlation_id": str(correlation_id), + "published_at": None, + }, + { + "value": 4, + "topic": "__out__", + "correlation_id": str(correlation_id), + "published_at": None, + }, ] # IN, OUT, one assert len(conn.topics) == 3 @@ -226,12 +256,37 @@ def test_invoke_join_then_subscribe(mocker: MockerFixture): # We get a single array result as chain_four waits for all publishers to finish # before operating on all elements published to topic_two as an array assert pubsub.invoke([2, 3], {"correlation_id": correlation_id}) == [26] - assert [{**m, "started_at": None} for m in conn.peek(correlation_id)] == [ - {"message": [2, 3], "topic_name": "__in__", "started_at": None}, - {"message": 12, "topic_name": "one", "started_at": None}, - {"message": 13, "topic_name": "one", "started_at": None}, - {"message": 25, "topic_name": "two", "started_at": None}, - {"message": 26, "topic_name": "__out__", "started_at": None}, + assert [{**m, "published_at": None} for m in conn.peek(correlation_id)] == [ + { + "value": [2, 3], + "topic": "__in__", + "correlation_id": str(correlation_id), + "published_at": None, + }, + { + "value": 12, + "topic": "one", + "correlation_id": str(correlation_id), + "published_at": None, + }, + { + "value": 13, + "topic": "one", + "correlation_id": str(correlation_id), + "published_at": None, + }, + { + "value": 25, + "topic": "two", + "correlation_id": str(correlation_id), + "published_at": None, + }, + { + "value": 26, + "topic": "__out__", + "correlation_id": str(correlation_id), + "published_at": None, + }, ] # After invoke returns the listeners were cleaned up