langgraph: add tests for remote graph interrupts (#4048)

This commit is contained in:
Vadym Barda
2025-03-27 15:17:29 +00:00
committed by GitHub
parent 7021ce3742
commit 96847e644b
+47 -4
View File
@@ -12,6 +12,7 @@ from langgraph_sdk.schema import StreamPart
from langgraph.errors import GraphInterrupt
from langgraph.pregel.remote import RemoteGraph
from langgraph.pregel.types import StateSnapshot
from langgraph.types import Interrupt
def test_with_config():
@@ -415,7 +416,19 @@ def test_stream():
StreamPart(event="values", data={"chunk": "data2"}),
StreamPart(event="values", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
StreamPart(event="updates", data={"__interrupt__": ()}),
StreamPart(
event="updates",
data={
"__interrupt__": [
{
"value": {"question": "Does this look good?"},
"resumable": True,
"ns": ["some_ns"],
"when": "during",
}
]
},
),
]
# call method / assertions
@@ -426,7 +439,7 @@ def test_stream():
# stream modes doesn't include 'updates'
stream_parts = []
with pytest.raises(GraphInterrupt):
with pytest.raises(GraphInterrupt) as exc:
for stream_part in remote_pregel.stream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
@@ -434,6 +447,15 @@ def test_stream():
):
stream_parts.append(stream_part)
assert exc.value.args[0] == [
Interrupt(
value={"question": "Does this look good?"},
resumable=True,
ns=["some_ns"],
when="during",
)
]
assert stream_parts == [
{"chunk": "data1"},
{"chunk": "data2"},
@@ -517,7 +539,19 @@ async def test_astream():
StreamPart(event="values", data={"chunk": "data2"}),
StreamPart(event="values", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
StreamPart(event="updates", data={"__interrupt__": ()}),
StreamPart(
event="updates",
data={
"__interrupt__": [
{
"value": {"question": "Does this look good?"},
"resumable": True,
"ns": ["some_ns"],
"when": "during",
}
]
},
),
]
mock_async_client.runs.stream.return_value = async_iter
@@ -529,7 +563,7 @@ async def test_astream():
# stream modes doesn't include 'updates'
stream_parts = []
with pytest.raises(GraphInterrupt):
with pytest.raises(GraphInterrupt) as exc:
async for stream_part in remote_pregel.astream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
@@ -537,6 +571,15 @@ async def test_astream():
):
stream_parts.append(stream_part)
assert exc.value.args[0] == [
Interrupt(
value={"question": "Does this look good?"},
resumable=True,
ns=["some_ns"],
when="during",
)
]
assert stream_parts == [
{"chunk": "data1"},
{"chunk": "data2"},