Update tests

This commit is contained in:
Nuno Campos
2025-01-22 16:30:43 -08:00
parent c43a9a4bd0
commit 23c73ae719
3 changed files with 43 additions and 39 deletions
+5 -1
View File
@@ -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")
+19 -19
View File
@@ -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": [],
},
@@ -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": [],
},