From 23c73ae7195cde40977c431e39da60f9ad59b2ff Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 22 Jan 2025 16:30:43 -0800 Subject: [PATCH] Update tests --- libs/langgraph/tests/test_pregel_async.py | 6 ++- libs/scheduler-kafka/tests/test_subgraph.py | 38 +++++++++---------- .../tests/test_subgraph_sync.py | 38 +++++++++---------- 3 files changed, 43 insertions(+), 39 deletions(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 166f8720a..fa1156859 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1612,7 +1612,11 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) + async def add_one_impl(x: int) -> int: + await asyncio.sleep(0.01 * x) + return x + 1 + + add_one = mocker.Mock(side_effect=add_one_impl) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = ( Channel.subscribe_to("inbox") diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index 2a6c9992a..934e3a614 100644 --- a/libs/scheduler-kafka/tests/test_subgraph.py +++ b/libs/scheduler-kafka/tests/test_subgraph.py @@ -15,7 +15,7 @@ from langgraph.graph.state import StateGraph from langgraph.pregel import Pregel from langgraph.scheduler.kafka import serde from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics -from tests.any import AnyDict, AnyInt +from tests.any import AnyDict from tests.drain import drain_topics_async from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage @@ -199,9 +199,9 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, + "subgraph_counter": None, + "call_counter": None, + "interrupt_counter": None, "null_resume": None, "resume": [], }, @@ -272,9 +272,9 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, + "subgraph_counter": None, + "call_counter": None, + "interrupt_counter": None, "null_resume": None, "resume": [], }, @@ -375,9 +375,9 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, + "subgraph_counter": None, + "call_counter": None, + "interrupt_counter": None, "null_resume": None, "resume": [], }, @@ -488,9 +488,9 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, + "subgraph_counter": None, + "call_counter": None, + "interrupt_counter": None, "null_resume": None, "resume": [], }, @@ -556,9 +556,9 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, + "subgraph_counter": None, + "call_counter": None, + "interrupt_counter": None, "null_resume": None, "resume": [], }, @@ -680,9 +680,9 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, + "subgraph_counter": None, + "call_counter": None, + "interrupt_counter": None, "null_resume": None, "resume": [], }, diff --git a/libs/scheduler-kafka/tests/test_subgraph_sync.py b/libs/scheduler-kafka/tests/test_subgraph_sync.py index c2c9a8fc1..50b33a39d 100644 --- a/libs/scheduler-kafka/tests/test_subgraph_sync.py +++ b/libs/scheduler-kafka/tests/test_subgraph_sync.py @@ -15,7 +15,7 @@ from langgraph.pregel import Pregel from langgraph.scheduler.kafka import serde from langgraph.scheduler.kafka.default_sync import DefaultProducer from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics -from tests.any import AnyDict, AnyInt +from tests.any import AnyDict from tests.drain import drain_topics from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage @@ -198,9 +198,9 @@ def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, + "subgraph_counter": None, + "call_counter": None, + "interrupt_counter": None, "null_resume": None, "resume": [], }, @@ -271,9 +271,9 @@ def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, + "subgraph_counter": None, + "call_counter": None, + "interrupt_counter": None, "null_resume": None, "resume": [], }, @@ -374,9 +374,9 @@ def test_subgraph_w_interrupt( "__pregel_task_id": history[0].tasks[0].id, "__pregel_previous": None, "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, + "subgraph_counter": None, + "call_counter": None, + "interrupt_counter": None, "null_resume": None, "resume": [], }, @@ -486,9 +486,9 @@ def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, + "subgraph_counter": None, + "call_counter": None, + "interrupt_counter": None, "null_resume": None, "resume": [], }, @@ -554,9 +554,9 @@ def test_subgraph_w_interrupt( "__pregel_previous": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, + "subgraph_counter": None, + "call_counter": None, + "interrupt_counter": None, "null_resume": None, "resume": [], }, @@ -678,9 +678,9 @@ def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { - "subgraph_counter": AnyInt(), - "call_counter": 0, - "interrupt_counter": -1, + "subgraph_counter": None, + "call_counter": None, + "interrupt_counter": None, "null_resume": None, "resume": [], },