Skip to content

Commit

Permalink
Update test
Browse files Browse the repository at this point in the history
  • Loading branch information
nfcampos committed Dec 10, 2024
1 parent 7cabc0a commit a7ac9ff
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 4 deletions.
8 changes: 4 additions & 4 deletions libs/langgraph/tests/test_pregel.py
Original file line number Diff line number Diff line change
Expand Up @@ -15016,16 +15016,16 @@ def test_command_goto_with_static_breakpoints(
class State(TypedDict):
"""The graph state."""

foo: str
foo: Annotated[str, operator.add]

def node1(state: State):
return {
"foo": state["foo"] + "|node-1",
"foo": "|node-1",
}

def node2(state: State):
return {
"foo": state["foo"] + "|node-2",
"foo": "|node-2",
}

builder = StateGraph(State)
Expand All @@ -15041,4 +15041,4 @@ def node2(state: State):
# Start the graph and interrupt at the first node
graph.invoke({"foo": "abc"}, config)
result = graph.invoke(Command(goto=["node2"]), config)
assert result == {"foo": "abc|node-2"}
assert result == {"foo": "abc|node-1|node-2|node-2"}
36 changes: 36 additions & 0 deletions libs/langgraph/tests/test_pregel_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -13295,3 +13295,39 @@ def step4(state: State):
],
"plan": [],
}


@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_command_goto_with_static_breakpoints(checkpointer_name: str) -> None:
"""Use Command goto with static breakpoints."""

class State(TypedDict):
"""The graph state."""

foo: Annotated[str, operator.add]

def node1(state: State):
return {
"foo": "|node-1",
}

def node2(state: State):
return {
"foo": "|node-2",
}

builder = StateGraph(State)
builder.add_node("node1", node1)
builder.add_node("node2", node2)
builder.add_edge(START, "node1")
builder.add_edge("node1", "node2")

async with awith_checkpointer(checkpointer_name) as checkpointer:
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"])

config = {"configurable": {"thread_id": str(uuid.uuid4())}}

# Start the graph and interrupt at the first node
await graph.ainvoke({"foo": "abc"}, config)
result = await graph.ainvoke(Command(goto=["node2"]), config)
assert result == {"foo": "abc|node-1|node-2|node-2"}

0 comments on commit a7ac9ff

Please sign in to comment.