diff --git a/examples/runnable-pregel.py b/examples/runnable-pregel.py index 5666fca29..8dc24f6b3 100644 --- a/examples/runnable-pregel.py +++ b/examples/runnable-pregel.py @@ -101,7 +101,9 @@ reviser_node = ( ) draft_revise_loop = Pregel( - (drafter_node, reviser_node, editor_node), + drafter_node, + reviser_node, + editor_node, input=question, output=draft, ) diff --git a/permchain/pregel.py b/permchain/pregel.py index 6e7a8ea55..7c6cd7e4e 100644 --- a/permchain/pregel.py +++ b/permchain/pregel.py @@ -237,8 +237,7 @@ class Pregel(Generic[Input, Output], RunnableSerializable[Input, Output]): def __init__( self, - processes: Sequence[PregelInvoke | PregelBatch], - *, + *processes: PregelInvoke | PregelBatch, input: Channel[Input, Any], output: Channel[Output, Any], step_timeout: Optional[float] = None, diff --git a/tests/test_invoke.py b/tests/test_invoke.py deleted file mode 100644 index aff324231..000000000 --- a/tests/test_invoke.py +++ /dev/null @@ -1,498 +0,0 @@ -from typing import Iterator -from uuid import uuid4 - -import pytest -from pytest_mock import MockerFixture - -from permchain.connection import PubSubMessage -from permchain.connection_inmemory import InMemoryPubSubConnection -from permchain.pubsub import PubSub -from permchain.topic import RunnableSubscriber, Topic - - -def clean_log( - logs: Iterator[PubSubMessage], correlation_id: bool | None = None -) -> list[PubSubMessage]: - if correlation_id is False: - return [{**m, "published_at": None, "correlation_id": None} for m in logs] - else: - return [{**m, "published_at": None} for m in logs] - - -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(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(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 == {} - - -@pytest.mark.skip("TODO") -def test_invoke_two_processes_in_out_interrupt(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(clear_on_disconnect=False) - pubsub_one = PubSub(chain_one, connection=conn) - pubsub_two = PubSub(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 both pubsubs, as a group - # The second picks up where the first left off - correlation_id = uuid4() - - # invoke() step 1 - assert clean_log(pubsub_one.stream(2, {"correlation_id": 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, one - assert len(conn.topics) == 2 - topic_one_full_name = conn.full_name(correlation_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 clean_log(pubsub_two.stream(None, {"correlation_id": correlation_id})) == [ - { - "value": None, - "topic": "__in__", - "correlation_id": str(correlation_id), - "published_at": None, - }, - { - "value": 4, - "topic": "__out__", - "correlation_id": str(correlation_id), - "published_at": None, - }, - ] - # listeners are still cleared, even though state is preserved - assert conn.listeners == {} - # 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 2 sentinel "end" values - assert queue.qsize() == 2 - - -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) - - for _ in range(10): - # 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_batch_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()) - - conn = InMemoryPubSubConnection() - pubsub = PubSub(processes=chains, connection=conn) - - for _ in range(10): - # Using in-memory conn internals to make assertions about pubsub - # If we start with 0 listeners - assert conn.listeners == {} - # Then invoke pubsub - assert pubsub.batch([2, 1, 3, 4, 5]) == [ - 2 + test_size, - 1 + test_size, - 3 + test_size, - 4 + test_size, - 5 + 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 only one of the two return values, as computation is closed - # as soon as we publish to OUT for the first time - assert pubsub.invoke(2) == 3 - - # After invoke returns the listeners were cleaned up - assert conn.listeners == {} - - -def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture): - add_one = mocker.Mock(side_effect=lambda x: x + 1) - add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x)) - topic_one = Topic("one") - topic_two = Topic("two") - chain_one = Topic.IN.subscribe() | add_one | topic_one.publish() - chain_two = topic_one.subscribe() | add_one | topic_two.publish() - chain_three = Topic.IN.subscribe() | add_one | topic_two.publish() - chain_four = topic_two.join() | add_10_each | Topic.OUT.publish() - - # Chains can be invoked directly for testing - assert chain_one.invoke(2) == 3 - assert chain_four.invoke([2, 3]) == [12, 13] - - conn = InMemoryPubSubConnection() - pubsub = PubSub((chain_one, chain_two, chain_three, chain_four), 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 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) == [13, 14] - - # After invoke returns the listeners were cleaned up - assert conn.listeners == {} - - -def test_invoke_join_then_subscribe(mocker: MockerFixture): - add_one = mocker.Mock(side_effect=lambda x: x + 1) - add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) - - topic_one = Topic("one") - topic_two = Topic("two") - - chain_one = Topic.IN.subscribe() | add_10_each | topic_one.publish_each() - chain_two = topic_one.join() | sum | topic_two.publish() - chain_three = topic_two.subscribe() | add_one | Topic.OUT.publish() - - # Chains can be invoked directly for testing - assert chain_two.invoke([2, 3]) == 5 - assert chain_three.invoke(5) == 6 - - correlation_id = uuid4() - conn = InMemoryPubSubConnection(clear_on_disconnect=False) - pubsub = PubSub((chain_one, chain_two, chain_three), 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 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 clean_log(pubsub.stream([2, 3], {"correlation_id": 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 - assert conn.listeners == {} - - -def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture): - conn = InMemoryPubSubConnection(clear_on_disconnect=False) - add_one = mocker.Mock(side_effect=lambda x: x + 1) - - inner_pubsub = PubSub( - (Topic.IN.subscribe() | add_one | Topic.OUT.publish(),), connection=conn - ) - - add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) - - topic_one = Topic("one") - topic_two = Topic("two") - - chain_one = Topic.IN.subscribe() | add_10_each | topic_one.publish_each() - chain_two = topic_one.join() | inner_pubsub.map() | sorted | topic_two.publish() - chain_three = topic_two.subscribe() | sum | Topic.OUT.publish() - - pubsub = PubSub((chain_one, chain_two, chain_three), connection=conn) - - # Using in-memory conn internals to make assertions about pubsub - # If we start with 0 listeners - assert conn.listeners == {} - - # Then invoke pubsub - for _ in range(10): - assert clean_log(pubsub.stream([2, 3]), correlation_id=False) == [ - { - "value": [2, 3], - "topic": "__in__", - "correlation_id": None, - "published_at": None, - }, - { - "value": 12, - "topic": "one", - "correlation_id": None, - "published_at": None, - }, - { - "value": 13, - "topic": "one", - "correlation_id": None, - "published_at": None, - }, - { - "value": [13, 14], - "topic": "two", - "correlation_id": None, - "published_at": None, - }, - { - "value": 27, - "topic": "__out__", - "correlation_id": None, - "published_at": None, - }, - ] - - # 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.OUT.publish() | 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 - # pubsub stopped executing after publishing to OUT, so only one value is returned - assert clean_log(pubsub.stream(2), correlation_id=False) == [ - { - "value": 2, - "topic": "__in__", - "correlation_id": None, - "published_at": None, - }, - { - "value": 3, - "topic": "__out__", - "correlation_id": None, - "published_at": None, - }, - ] - - # 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) is None - - # 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) is None - - # After invoke returns the listeners were cleaned up - assert conn.listeners == {} - - -@pytest.mark.skip("TODO") -def test_invoke_two_processes_simple_cycle(mocker: MockerFixture) -> None: - 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] diff --git a/tests/test_pregel.py b/tests/test_pregel.py new file mode 100644 index 000000000..a0e8104ee --- /dev/null +++ b/tests/test_pregel.py @@ -0,0 +1,250 @@ +import time +from uuid import uuid4 + +import pytest +from pytest_mock import MockerFixture +from langchain.schema.runnable import RunnablePassthrough + +from permchain import Pregel, channels +from permchain.pregel import PregelInvoke + + +def test_invoke_single_process_in_out(mocker: MockerFixture): + input = channels.LastValue[int]("input") + output = channels.LastValue[int]("output") + + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output) + + pubsub = Pregel(chain, input=input, output=output) + + # Then invoke pubsub + assert pubsub.invoke(2) == 3 + + +def test_invoke_two_processes_in_out(mocker: MockerFixture): + input = channels.LastValue[int]("input") + output = channels.LastValue[int]("output") + inbox = channels.Inbox[int]("inbox") + + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(inbox) + chain_two = Pregel.subscribe_to_each(inbox) | add_one | Pregel.send_to(output) + + pubsub = Pregel(chain_one, chain_two, input=input, output=output) + + # Then invoke pubsub + assert pubsub.invoke(2) == 4 + + +def test_batch_two_processes_in_out(mocker: MockerFixture): + def add_one_with_delay(inp: int) -> int: + time.sleep(inp / 10) + return inp + 1 + + input = channels.LastValue[int]("input") + output = channels.LastValue[int]("output") + one = channels.LastValue[int]("one") + chain_one = Pregel.subscribe_to(input) | add_one_with_delay | Pregel.send_to(one) + chain_two = Pregel.subscribe_to(one) | add_one_with_delay | Pregel.send_to(output) + + pubsub = Pregel(chain_one, chain_two, input=input, output=output) + + # Then invoke pubsub + assert pubsub.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] + + +def test_invoke_many_processes_in_out(mocker: MockerFixture): + test_size = 100 + + input = channels.LastValue[int]("input") + output = channels.LastValue[int]("output") + topics: list[channels.Channel] = [channels.LastValue[int]("zero")] + + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chains: list[PregelInvoke] = [ + Pregel.subscribe_to(input) | add_one | Pregel.send_to(topics[0]) + ] + for i in range(test_size - 2): + topics.append(channels.LastValue[int](str(i))) + chains.append( + Pregel.subscribe_to(topics[-2]) | add_one | Pregel.send_to(topics[-1]) + ) + chains.append(Pregel.subscribe_to(topics[-1]) | add_one | Pregel.send_to(output)) + + pubsub = Pregel(*chains, input=input, output=output) + + for _ in range(10): + # Then invoke pubsub + assert pubsub.invoke(2, {"recursion_limit": test_size}) == 2 + test_size + + +def test_batch_many_processes_in_out(mocker: MockerFixture): + test_size = 100 + + input = channels.LastValue[int]("input") + output = channels.LastValue[int]("output") + topics: list[channels.Channel] = [channels.LastValue[int]("zero")] + + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chains: list[PregelInvoke] = [ + Pregel.subscribe_to(input) | add_one | Pregel.send_to(topics[0]) + ] + for i in range(test_size - 2): + topics.append(channels.LastValue[int](str(i))) + chains.append( + Pregel.subscribe_to(topics[-2]) | add_one | Pregel.send_to(topics[-1]) + ) + chains.append(Pregel.subscribe_to(topics[-1]) | add_one | Pregel.send_to(output)) + + pubsub = Pregel(*chains, input=input, output=output) + + for _ in range(10): + # Then invoke pubsub + assert pubsub.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [ + 2 + test_size, + 1 + test_size, + 3 + test_size, + 4 + test_size, + 5 + test_size, + ] + + +def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture): + add_one = mocker.Mock(side_effect=lambda x: x + 1) + + input = channels.LastValue[int]("input") + output = channels.LastValue[int]("output") + + chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output) + chain_two = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output) + + pubsub = Pregel(chain_one, chain_two, input=input, output=output) + + with pytest.raises(channels.InvalidUpdateError): + # LastValue channels can only be updated once per iteration + pubsub.invoke(2) + + +def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture): + add_one = mocker.Mock(side_effect=lambda x: x + 1) + + input = channels.LastValue[int]("input") + output = channels.Inbox[int]("output") + + chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output) + chain_two = Pregel.subscribe_to(input) | add_one | Pregel.send_to(output) + + pubsub = Pregel(chain_one, chain_two, input=input, output=output) + + # An Inbox channel accumulates updates into a sequence + assert pubsub.invoke(2) == (3, 3) + + +def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture): + add_one = mocker.Mock(side_effect=lambda x: x + 1) + add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x)) + + input = channels.LastValue[int]("input") + output = channels.LastValue[int]("output") + inbox = channels.Inbox[int]("inbox") + + chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(inbox) + chain_three = Pregel.subscribe_to(input) | add_one | Pregel.send_to(inbox) + chain_four = Pregel.subscribe_to(inbox) | add_10_each | Pregel.send_to(output) + + pubsub = Pregel(chain_one, chain_three, chain_four, input=input, output=output) + + # Then invoke pubsub + # 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 + for _ in range(100): + assert pubsub.invoke(2) == [13, 13] + + +def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture): + add_one = mocker.Mock(side_effect=lambda x: x + 1) + add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) + + input = channels.LastValue[int]("input") + output = channels.LastValue[int]("output") + + inner_pubsub = Pregel( + Pregel.subscribe_to(input) | add_one | Pregel.send_to(output), + input=input, + output=output, + ) + + inbox_one = channels.Inbox[int]("inbox_one") + outbox_one = channels.LastValue[int]("outbox_one") + + chain_one = ( + Pregel.subscribe_to(input) | add_10_each | Pregel.send_to(inbox_one).map() + ) + chain_two = ( + Pregel.subscribe_to(inbox_one) + | inner_pubsub.map() + | sorted + | Pregel.send_to(outbox_one) + ) + chain_three = Pregel.subscribe_to(outbox_one) | sum | Pregel.send_to(output) + + pubsub = Pregel(chain_one, chain_two, chain_three, input=input, output=output) + + # Then invoke pubsub + for _ in range(10): + assert pubsub.invoke([2, 3]) == 27 + + +def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture): + add_one = mocker.Mock(side_effect=lambda x: x + 1) + + input = channels.LastValue[int]("input") + output = channels.LastValue[int]("output") + between = channels.LastValue[int]("between") + + chain_one = ( + Pregel.subscribe_to(input) + | add_one + | Pregel.send_to( + {output: RunnablePassthrough(), between: RunnablePassthrough()} + ) + ) + chain_two = Pregel.subscribe_to(between) | add_one | Pregel.send_to(output) + + pubsub = Pregel(chain_one, chain_two, input=input, output=output) + + # Then invoke pubsub + assert [c for c in pubsub.stream(2)] == [3, 4] + + +def test_invoke_two_processes_no_out(mocker: MockerFixture): + input = channels.LastValue[int]("input") + output = channels.LastValue[int]("output") + between = channels.LastValue[int]("between") + + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain_one = Pregel.subscribe_to(input) | add_one | Pregel.send_to(between) + chain_two = Pregel.subscribe_to(between) | add_one + + pubsub = Pregel(chain_one, chain_two, input=input, output=output) + + # 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) is None + + +def test_invoke_two_processes_no_in(mocker: MockerFixture): + input = channels.LastValue[int]("input") + output = channels.LastValue[int]("output") + between = channels.LastValue[int]("between") + + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain_one = Pregel.subscribe_to(between) | add_one | Pregel.send_to(output) + chain_two = Pregel.subscribe_to(between) | add_one + + pubsub = Pregel(chain_one, chain_two, input=input, output=output) + + with pytest.raises(ValueError): + assert pubsub.invoke(2) is None