Add first tests, fix some minor issues

- Temporarily comment out code using future runnables features
2026-09-12 04:37:51 +02:00 · 2023-08-29 14:10:52 +01:00
parent af070018c2
commit 7f3eae759f
5 changed files with 315 additions and 39 deletions
@@ -1,27 +1,31 @@
 import threading
+import queue
 from collections import defaultdict
 from typing import Any

+q = queue.Queue()
+
 from permchain.connection import PubSubConnection, PubSubListener


 class InMemoryPubSubConnection(PubSubConnection):
-    topics: defaultdict[str, list[PubSubListener]]
-    lock: threading.Lock
+    listeners: defaultdict[str, list[PubSubListener]]
+    lock: threading.RLock

    def __init__(self) -> None:
-        self.topics = defaultdict(list)
-        self.lock = threading.Lock()
+        self.listeners = defaultdict(list)
+        self.lock = threading.RLock()

    def listen(self, topic_name: str, listener: PubSubListener) -> None:
        with self.lock:
-            self.topics[topic_name].append(listener)
+            self.listeners[topic_name].append(listener)

    def send(self, topic_name: str, message: Any) -> None:
        with self.lock:
-            for listener in self.topics[topic_name]:
+            for listener in self.listeners[topic_name]:
                listener(message)

    def disconnect(self, topic_name: str) -> None:
        with self.lock:
-            self.topics[topic_name] = []
+            if topic_name in self.listeners:
+                del self.listeners[topic_name]
@@ -1,6 +1,7 @@
 from __future__ import annotations

 import queue
+import threading
 from abc import ABC
 from concurrent.futures import CancelledError, Future
 from functools import partial
@@ -22,7 +23,6 @@ from langchain.load.serializable import Serializable
 from langchain.schema.runnable import Runnable, RunnableConfig, patch_config
 from langchain.schema.runnable.base import Runnable
 from langchain.schema.runnable.config import get_executor_for_config
-from tenacity import BaseRetrying

 from permchain.connection import PubSubConnection
 from permchain.constants import CONFIG_GET_KEY, CONFIG_SEND_KEY
@@ -80,6 +80,11 @@ class PubSub(Serializable, Runnable[Any, Any], ABC):
            lambda r: r.topic.name == INPUT_TOPIC, self.processes
        )

+        if not input_processes:
+            # TODO raise exception?
+            return
+
+        # Consume input iterator into a single value
        input_value = None
        for chunk in input:
            if input_value is None:
@@ -106,28 +111,30 @@ class PubSub(Serializable, Runnable[Any, Any], ABC):
                else:
                    self.connection.send(prefix_topic_name(topic_name), message)

-            def cleanup_run(fut: Future) -> None:
-                """Cleanup after a process runs."""
-                inflight.discard(fut)
-
-                try:
-                    exc = fut.exception()
-                except CancelledError:
-                    exc = None
-                except Exception as e:
-                    exc = e
-                if exc is not None:
-                    exceptions.append(exc)
-
-                # Close output iterator if
-                # - all processes are done, or
-                # - an exception occurred
-                if not inflight or exc is not None:
-                    output.close()
-
            def run_once(process: RunnableSubscriber[Any], value: Any) -> None:
                """Run a process once."""

+                def cleanup_run(fut: Future) -> None:
+                    """Cleanup after a process runs."""
+                    inflight.discard(fut)
+
+                    try:
+                        exc = fut.exception()
+                    except CancelledError:
+                        exc = None
+                    except Exception as e:
+                        exc = e
+                    if exc is not None:
+                        exceptions.append(exc)
+
+                    # Close output iterator if
+                    # - all processes are done, or
+                    # - an exception occurred
+                    if not inflight or exc is not None:
+                        # TODO remove timer once the condition includes
+                        # checking for messages not yet read from the topics
+                        threading.Timer(0.1, output.close).start()
+
                def get(topic_name: str) -> Any:
                    if topic_name == INPUT_TOPIC:
                        return input_value
@@ -146,7 +153,7 @@ class PubSub(Serializable, Runnable[Any, Any], ABC):
                        **patch_config(
                            config,
                            callbacks=run_manager.get_child(),
-                            run_name=f"Topic: {process.topic.name}",
+                            # run_name=f"Topic: {process.topic.name}",
                        ),
                        CONFIG_SEND_KEY: send,
                        CONFIG_GET_KEY: get,
@@ -173,8 +180,8 @@ class PubSub(Serializable, Runnable[Any, Any], ABC):
                    yield chunk
            finally:
                # Disconnect from all topics
-                for process in listener_processes:
-                    self.connection.disconnect(prefix_topic_name(process.topic.name))
+                for topic_name in set(p.topic.name for p in listener_processes):
+                    self.connection.disconnect(prefix_topic_name(topic_name))

                # Cancel all inflight futures
                while inflight:
@@ -761,6 +761,16 @@ files = [
    {file = "defusedxml-0.7.1.tar.gz", hash = "sha256:1bb3032db185915b62d7c6209c5a8792be6a32ab2fedacc84e01b52c51aa3e69"},
 ]

+[[package]]
+name = "docopt"
+version = "0.6.2"
+description = "Pythonic argument parser, that will make you smile"
+optional = false
+python-versions = "*"
+files = [
+    {file = "docopt-0.6.2.tar.gz", hash = "sha256:49b3a825280bd66b3aa83585ef59c4a8c82f2c8a522dbe754a8bc8d08c85c491"},
+]
+
 [[package]]
 name = "exceptiongroup"
 version = "1.1.2"
@@ -2477,6 +2487,22 @@ pytest = ">=5.0"
 [package.extras]
 dev = ["pre-commit", "pytest-asyncio", "tox"]

+[[package]]
+name = "pytest-watch"
+version = "4.2.0"
+description = "Local continuous test runner with pytest and watchdog."
+optional = false
+python-versions = "*"
+files = [
+    {file = "pytest-watch-4.2.0.tar.gz", hash = "sha256:06136f03d5b361718b8d0d234042f7b2f203910d8568f63df2f866b547b3d4b9"},
+]
+
+[package.dependencies]
+colorama = ">=0.3.3"
+docopt = ">=0.4.0"
+pytest = ">=2.6.4"
+watchdog = ">=0.6.0"
+
 [[package]]
 name = "python-dateutil"
 version = "2.8.2"
@@ -3301,6 +3327,45 @@ secure = ["certifi", "cryptography (>=1.9)", "idna (>=2.0.0)", "pyopenssl (>=17.
 socks = ["pysocks (>=1.5.6,!=1.5.7,<2.0)"]
 zstd = ["zstandard (>=0.18.0)"]

+[[package]]
+name = "watchdog"
+version = "3.0.0"
+description = "Filesystem events monitoring"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "watchdog-3.0.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:336adfc6f5cc4e037d52db31194f7581ff744b67382eb6021c868322e32eef41"},
+    {file = "watchdog-3.0.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:a70a8dcde91be523c35b2bf96196edc5730edb347e374c7de7cd20c43ed95397"},
+    {file = "watchdog-3.0.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:adfdeab2da79ea2f76f87eb42a3ab1966a5313e5a69a0213a3cc06ef692b0e96"},
+    {file = "watchdog-3.0.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:2b57a1e730af3156d13b7fdddfc23dea6487fceca29fc75c5a868beed29177ae"},
+    {file = "watchdog-3.0.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:7ade88d0d778b1b222adebcc0927428f883db07017618a5e684fd03b83342bd9"},
+    {file = "watchdog-3.0.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:7e447d172af52ad204d19982739aa2346245cc5ba6f579d16dac4bfec226d2e7"},
+    {file = "watchdog-3.0.0-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:9fac43a7466eb73e64a9940ac9ed6369baa39b3bf221ae23493a9ec4d0022674"},
+    {file = "watchdog-3.0.0-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:8ae9cda41fa114e28faf86cb137d751a17ffd0316d1c34ccf2235e8a84365c7f"},
+    {file = "watchdog-3.0.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:25f70b4aa53bd743729c7475d7ec41093a580528b100e9a8c5b5efe8899592fc"},
+    {file = "watchdog-3.0.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:4f94069eb16657d2c6faada4624c39464f65c05606af50bb7902e036e3219be3"},
+    {file = "watchdog-3.0.0-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:7c5f84b5194c24dd573fa6472685b2a27cc5a17fe5f7b6fd40345378ca6812e3"},
+    {file = "watchdog-3.0.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:3aa7f6a12e831ddfe78cdd4f8996af9cf334fd6346531b16cec61c3b3c0d8da0"},
+    {file = "watchdog-3.0.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:233b5817932685d39a7896b1090353fc8efc1ef99c9c054e46c8002561252fb8"},
+    {file = "watchdog-3.0.0-pp37-pypy37_pp73-macosx_10_9_x86_64.whl", hash = "sha256:13bbbb462ee42ec3c5723e1205be8ced776f05b100e4737518c67c8325cf6100"},
+    {file = "watchdog-3.0.0-pp38-pypy38_pp73-macosx_10_9_x86_64.whl", hash = "sha256:8f3ceecd20d71067c7fd4c9e832d4e22584318983cabc013dbf3f70ea95de346"},
+    {file = "watchdog-3.0.0-pp39-pypy39_pp73-macosx_10_9_x86_64.whl", hash = "sha256:c9d8c8ec7efb887333cf71e328e39cffbf771d8f8f95d308ea4125bf5f90ba64"},
+    {file = "watchdog-3.0.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:0e06ab8858a76e1219e68c7573dfeba9dd1c0219476c5a44d5333b01d7e1743a"},
+    {file = "watchdog-3.0.0-py3-none-manylinux2014_armv7l.whl", hash = "sha256:d00e6be486affb5781468457b21a6cbe848c33ef43f9ea4a73b4882e5f188a44"},
+    {file = "watchdog-3.0.0-py3-none-manylinux2014_i686.whl", hash = "sha256:c07253088265c363d1ddf4b3cdb808d59a0468ecd017770ed716991620b8f77a"},
+    {file = "watchdog-3.0.0-py3-none-manylinux2014_ppc64.whl", hash = "sha256:5113334cf8cf0ac8cd45e1f8309a603291b614191c9add34d33075727a967709"},
+    {file = "watchdog-3.0.0-py3-none-manylinux2014_ppc64le.whl", hash = "sha256:51f90f73b4697bac9c9a78394c3acbbd331ccd3655c11be1a15ae6fe289a8c83"},
+    {file = "watchdog-3.0.0-py3-none-manylinux2014_s390x.whl", hash = "sha256:ba07e92756c97e3aca0912b5cbc4e5ad802f4557212788e72a72a47ff376950d"},
+    {file = "watchdog-3.0.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:d429c2430c93b7903914e4db9a966c7f2b068dd2ebdd2fa9b9ce094c7d459f33"},
+    {file = "watchdog-3.0.0-py3-none-win32.whl", hash = "sha256:3ed7c71a9dccfe838c2f0b6314ed0d9b22e77d268c67e015450a29036a81f60f"},
+    {file = "watchdog-3.0.0-py3-none-win_amd64.whl", hash = "sha256:4c9956d27be0bb08fc5f30d9d0179a855436e655f046d288e2bcc11adfae893c"},
+    {file = "watchdog-3.0.0-py3-none-win_ia64.whl", hash = "sha256:5d9f3a10e02d7371cd929b5d8f11e87d4bad890212ed3901f9b4d68767bee759"},
+    {file = "watchdog-3.0.0.tar.gz", hash = "sha256:4d98a320595da7a7c5a18fc48cb633c2e73cda78f93cac2ef42d42bf609a33f9"},
+]
+
+[package.extras]
+watchmedo = ["PyYAML (>=3.10)"]
+
 [[package]]
 name = "wcwidth"
 version = "0.2.6"
@@ -3470,4 +3535,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.8.1,<4.0"
-content-hash = "efcc0b5e1b923cc732d3f214a48e9e5b362ceb415839307fa82ac91839ac5aa7"
+content-hash = "4ab612eade28e540c067984413cd631684426bb7f51b389cd27d0c41374829d1"
@@ -22,6 +22,7 @@ pytest-dotenv = "^0.5.2"
 pytest-asyncio = "^0.20.3"
 pytest-mock  = "^3.10.0"
 syrupy = "^4.0.2"
+pytest-watch = "^4.2.0"

 [tool.poetry.group.lint.dependencies]
 ruff = "^0.0.249"
@@ -40,11 +41,7 @@ setuptools = "^67.6.1"
 openai = "^0.27.8"

 [tool.ruff]
-select = [
-  "E",  # pycodestyle
-  "F",  # pyflakes
-  "I",  # isort
-]
+select = [ "E", "F", "I" ]

 [tool.mypy]
 ignore_missing_imports = "True"
@@ -52,9 +49,7 @@ disallow_untyped_defs = "True"
 exclude = ["notebooks", "examples", "example_data"]

 [tool.coverage.run]
-omit = [
-    "tests/*",
-]
+omit = ["tests/*"]

 [build-system]
 requires = ["poetry-core>=1.0.0"]
@@ -0,0 +1,205 @@
+import pytest
+from pytest_mock import MockerFixture
+from permchain.connection_inmemory import InMemoryPubSubConnection
+from permchain.pubsub import PubSub
+
+from permchain.topic import RunnableSubscriber, Topic
+
+
+def test_invoke_single_process_in_out(mocker: MockerFixture):
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain = Topic.IN.subscribe() | add_one | Topic.OUT.publish()
+
+    # Chains can be invoked directly for testing
+    assert chain.invoke(2) == 3
+
+    conn = InMemoryPubSubConnection()
+    pubsub = PubSub(processes=(chain,), connection=conn)
+
+    # Using in-memory conn internals to make assertions about pubsub
+    # If we start with 0 listeners
+    assert conn.listeners == {}
+    # Then invoke pubsub
+    assert pubsub.invoke(2) == [3]
+    # After invoke returns the listeners were cleaned up
+    assert conn.listeners == {}
+
+
+def test_invoke_two_processes_in_out(mocker: MockerFixture):
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    topic_one = Topic("one")
+    chain_one = Topic.IN.subscribe() | add_one | topic_one.publish()
+    chain_two = topic_one.subscribe() | add_one | Topic.OUT.publish()
+
+    # Chains can be invoked directly for testing
+    assert chain_one.invoke(2) == 3
+    assert chain_two.invoke(2) == 3
+
+    conn = InMemoryPubSubConnection()
+    pubsub = PubSub(processes=(chain_one, chain_two), connection=conn)
+
+    # Using in-memory conn internals to make assertions about pubsub
+    # If we start with 0 listeners
+    assert conn.listeners == {}
+    # Then invoke pubsub
+    assert pubsub.invoke(2) == [4]
+    # After invoke returns the listeners were cleaned up
+    assert conn.listeners == {}
+
+
+def test_invoke_many_processes_in_out(mocker: MockerFixture):
+    test_size = 100
+
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    topics: list[Topic] = [Topic("zero")]
+    chains: list[RunnableSubscriber] = [
+        Topic.IN.subscribe() | add_one | topics[0].publish()
+    ]
+    for i in range(test_size - 2):
+        topics.append(Topic(str(i)))
+        chains.append(topics[-2].subscribe() | add_one | topics[-1].publish())
+    chains.append(topics[-1].subscribe() | add_one | Topic.OUT.publish())
+
+    # Chains can be invoked directly for testing
+    for chain in chains:
+        assert chain.invoke(2) == 3
+
+    conn = InMemoryPubSubConnection()
+    pubsub = PubSub(processes=chains, connection=conn)
+
+    # Using in-memory conn internals to make assertions about pubsub
+    # If we start with 0 listeners
+    assert conn.listeners == {}
+    # Then invoke pubsub
+    assert pubsub.invoke(2) == [2 + test_size]
+    # After invoke returns the listeners were cleaned up
+    assert conn.listeners == {}
+
+
+def test_invoke_two_processes_two_in_two_out(mocker: MockerFixture):
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    chain_one = Topic.IN.subscribe() | add_one | Topic.OUT.publish()
+    chain_two = Topic.IN.subscribe() | add_one | Topic.OUT.publish()
+
+    # Chains can be invoked directly for testing
+    assert chain_one.invoke(2) == 3
+    assert chain_two.invoke(2) == 3
+
+    conn = InMemoryPubSubConnection()
+    pubsub = PubSub(processes=(chain_one, chain_two), connection=conn)
+
+    # Using in-memory conn internals to make assertions about pubsub
+    # If we start with 0 listeners
+    assert conn.listeners == {}
+
+    # Then invoke pubsub
+    # We get two equal results as the two chains do the same thing
+    assert pubsub.invoke(2) == [3, 3]
+
+    # After invoke returns the listeners were cleaned up
+    assert conn.listeners == {}
+
+
+def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture):
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    topic_one = Topic("one")
+    # Topic.publish() is passthrough so we can publish to multiple topics in sequence
+    chain_one = (
+        Topic.IN.subscribe() | add_one | topic_one.publish() | Topic.OUT.publish()
+    )
+    chain_two = topic_one.subscribe() | add_one | Topic.OUT.publish()
+
+    # Chains can be invoked directly for testing
+    assert chain_one.invoke(2) == 3
+    assert chain_two.invoke(2) == 3
+
+    conn = InMemoryPubSubConnection()
+    pubsub = PubSub(processes=(chain_one, chain_two), connection=conn)
+
+    # Using in-memory conn internals to make assertions about pubsub
+    # If we start with 0 listeners
+    assert conn.listeners == {}
+
+    # Then invoke pubsub
+    # pubsub didn't stop executing after getting the first return value
+    # the values arrive in the order they are produced
+    assert pubsub.invoke(2) == [3, 4]
+
+    # After invoke returns the listeners were cleaned up
+    assert conn.listeners == {}
+
+
+def test_invoke_two_processes_no_out(mocker: MockerFixture):
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    topic_one = Topic("one")
+    chain_one = Topic.IN.subscribe() | add_one | topic_one.publish()
+    chain_two = topic_one.subscribe() | add_one
+
+    # Chains can be invoked directly for testing
+    assert chain_one.invoke(2) == 3
+    assert chain_two.invoke(2) == 3
+
+    conn = InMemoryPubSubConnection()
+    pubsub = PubSub(processes=(chain_one, chain_two), connection=conn)
+
+    # Using in-memory conn internals to make assertions about pubsub
+    # If we start with 0 listeners
+    assert conn.listeners == {}
+
+    # Then invoke pubsub
+    # It finishes executing (once no more messages being published)
+    # but returns nothing, as nothing was published to OUT topic
+    assert pubsub.invoke(2) == []
+
+    # After invoke returns the listeners were cleaned up
+    assert conn.listeners == {}
+
+
+def test_invoke_two_processes_no_in(mocker: MockerFixture):
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    topic_one = Topic("one")
+    chain_one = topic_one.subscribe() | add_one | Topic.OUT.publish()
+    chain_two = topic_one.subscribe() | add_one | Topic.OUT.publish()
+
+    # Chains can be invoked directly for testing
+    assert chain_one.invoke(2) == 3
+    assert chain_two.invoke(2) == 3
+
+    conn = InMemoryPubSubConnection()
+    pubsub = PubSub(processes=(chain_one, chain_two), connection=conn)
+
+    # Using in-memory conn internals to make assertions about pubsub
+    # If we start with 0 listeners
+    assert conn.listeners == {}
+
+    # Then invoke pubsub
+    # It returns without any output as there is nothing to run
+    assert pubsub.invoke(2) == []
+
+    # After invoke returns the listeners were cleaned up
+    assert conn.listeners == {}
+
+
+@pytest.mark.skip("TODO")
+def test_invoke_two_processes_simple_cycle(mocker: MockerFixture):
+    add_one = mocker.Mock(side_effect=lambda x: x + 1)
+    topic_one = Topic("one")
+    chain_one = Topic.IN.subscribe() | add_one | topic_one.publish()
+    chain_two = topic_one.subscribe() | add_one | topic_one.publish()
+
+    # Chains can be invoked directly for testing
+    assert chain_one.invoke(2) == 3
+    assert chain_two.invoke(2) == 3
+
+    conn = InMemoryPubSubConnection()
+    pubsub = PubSub(processes=(chain_one, chain_two), connection=conn)
+
+    # Using in-memory conn internals to make assertions about pubsub
+    # If we start with 0 listeners
+    assert conn.listeners == {}
+    # Then invoke pubsub
+    with pytest.raises(RecursionError):
+        pubsub.invoke(2)
+    # After invoke returns the listeners were cleaned up
+    for key in conn.listeners:
+        assert not conn.listeners[key]