Skip to content

Commit

Permalink
Update tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@nfcampos
nfcampos committed Jan 23, 2025
1 parent c43a9a4 commit 23c73ae
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 39 deletions.
6 changes: 5 additions & 1 deletion libs/langgraph/tests/test_pregel_async.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
38 changes: 19 additions & 19 deletions libs/scheduler-kafka/tests/test_subgraph.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
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

Expand Down Expand Up @@ -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": [],
},
Expand Down Expand Up @@ -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": [],
},
Expand Down Expand Up @@ -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": [],
},
Expand Down Expand Up @@ -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": [],
},
Expand Down Expand Up @@ -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": [],
},
Expand Down Expand Up @@ -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": [],
},
Expand Down
38 changes: 19 additions & 19 deletions libs/scheduler-kafka/tests/test_subgraph_sync.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
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

Expand Down Expand Up @@ -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": [],
},
Expand Down Expand Up @@ -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": [],
},
Expand Down Expand Up @@ -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": [],
},
Expand Down Expand Up @@ -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": [],
},
Expand Down Expand Up @@ -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": [],
},
Expand Down Expand Up @@ -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": [],
},
Expand Down

0 comments on commit 23c73ae

Please sign in to comment.