From a6c6b8211fa71a51e5b5e64227f5854fa345922f Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 8 May 2024 10:43:31 -0700
Subject: [PATCH] Add tests for calling update on an empty thread

- this would eg be what you'd do in an api like openai assistants api
---
 langgraph/pregel/__init__.py |  6 ++----
 tests/test_pregel.py         | 36 ++++++++++++++++++++++++++++++++
 tests/test_pregel_async.py   | 40 ++++++++++++++++++++++++++++++++++++
 3 files changed, 78 insertions(+), 4 deletions(-)

diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py
index db65432ec..ebc8197f5 100644
--- a/langgraph/pregel/__init__.py
+++ b/langgraph/pregel/__init__.py
@@ -485,9 +485,7 @@ class Pregel(
                 create_checkpoint(checkpoint, channels),
                 {
                     "source": "update",
-                    "step": saved.metadata.get("step", 0) + 1
-                    if saved.metadata
-                    else None,
+                    "step": saved.metadata.get("step", 0) + 1 if saved else 0,
                 },
             )
 
@@ -560,7 +558,7 @@ class Pregel(
                 create_checkpoint(checkpoint, channels),
                 {
                     "source": "update",
-                    "step": saved.metadata.get("step", 0) + 1 if saved else None,
+                    "step": saved.metadata.get("step", 0) + 1 if saved else 0,
                 },
             )
 
diff --git a/tests/test_pregel.py b/tests/test_pregel.py
index 47dc14e79..46c59ad81 100644
--- a/tests/test_pregel.py
+++ b/tests/test_pregel.py
@@ -3881,6 +3881,42 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
             parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config,
         )
 
+        thread3 = {"configurable": {"thread_id": "3"}}
+        # update an empty thread before first run
+        uconfig = tool_two.update_state(thread3, {"my_key": "key", "market": "DE"})
+        # check current state
+        assert tool_two.get_state(thread3) == StateSnapshot(
+            values={"my_key": "key", "market": "DE"},
+            next=("prepare",),
+            config=uconfig,
+            metadata={"source": "update", "step": 0},
+        )
+        # run from this point
+        assert tool_two.invoke(None, thread3) == {
+            "my_key": "key prepared",
+            "market": "DE",
+        }
+        # get state after first node
+        assert tool_two.get_state(thread3) == StateSnapshot(
+            values={"my_key": "key prepared", "market": "DE"},
+            next=("tool_two_slow",),
+            config=tool_two.checkpointer.get_tuple(thread3).config,
+            metadata={"source": "loop", "step": 1},
+            parent_config=uconfig,
+        )
+        # resume, for same result as above
+        assert tool_two.invoke(None, thread3, debug=1) == {
+            "my_key": "key prepared slow finished",
+            "market": "DE",
+        }
+        assert tool_two.get_state(thread3) == StateSnapshot(
+            values={"my_key": "key prepared slow finished", "market": "DE"},
+            next=(),
+            config=tool_two.checkpointer.get_tuple(thread3).config,
+            metadata={"source": "loop", "step": 3},
+            parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config,
+        )
+
 
 def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> None:
     def sorted_add(
diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py
index 8d5b3b355..51ecd43e8 100644
--- a/tests/test_pregel_async.py
+++ b/tests/test_pregel_async.py
@@ -3348,6 +3348,46 @@ async def test_branch_then() -> None:
             ][-1].config,
         )
 
+        thread3 = {"configurable": {"thread_id": "3"}}
+        # update an empty thread before first run
+        uconfig = await tool_two.aupdate_state(
+            thread3, {"my_key": "key", "market": "DE"}
+        )
+        # check current state
+        assert await tool_two.aget_state(thread3) == StateSnapshot(
+            values={"my_key": "key", "market": "DE"},
+            next=("prepare",),
+            config=uconfig,
+            metadata={"source": "update", "step": 0},
+        )
+        # run from this point
+        assert await tool_two.ainvoke(None, thread3) == {
+            "my_key": "key prepared",
+            "market": "DE",
+        }
+        # get state after first node
+        assert await tool_two.aget_state(thread3) == StateSnapshot(
+            values={"my_key": "key prepared", "market": "DE"},
+            next=("tool_two_slow",),
+            config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
+            metadata={"source": "loop", "step": 1},
+            parent_config=uconfig,
+        )
+        # resume, for same result as above
+        assert await tool_two.ainvoke(None, thread3, debug=1) == {
+            "my_key": "key prepared slow finished",
+            "market": "DE",
+        }
+        assert await tool_two.aget_state(thread3) == StateSnapshot(
+            values={"my_key": "key prepared slow finished", "market": "DE"},
+            next=(),
+            config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
+            metadata={"source": "loop", "step": 3},
+            parent_config=[
+                c async for c in tool_two.checkpointer.alist(thread3, limit=2)
+            ][-1].config,
+        )
+
 
 async def test_in_one_fan_out_state_graph_waiting_edge() -> None:
     def sorted_add(