From a7bb96da98c0086da7bc9c171253d884882f5327 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Wed, 15 Jan 2025 10:21:40 -0800
Subject: [PATCH] Fix

---
 libs/langgraph/tests/test_pregel_async.py | 45 +++++++++++++----------
 1 file changed, 25 insertions(+), 20 deletions(-)

diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index d826daf99..b64aae8a3 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -6714,12 +6714,15 @@ async def test_falsy_return_from_task() -> None:
     await graph.ainvoke(Command(resume="123"), configurable)
 
 
-async def test_multiple_interrupts_imperative() -> None:
+@pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="Python 3.11+ is required for async contextvars support",
+)
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_multiple_interrupts_imperative(checkpointer_name: str) -> None:
     """Test multiple interrupts with an imperative API."""
-    from langgraph.checkpoint.memory import MemorySaver
     from langgraph.func import entrypoint, task
 
-    checkpointer = MemorySaver()
     counter = 0
 
     @task
@@ -6729,24 +6732,26 @@ async def test_multiple_interrupts_imperative() -> None:
         counter += 1
         return 2 * x
 
-    @entrypoint(checkpointer=checkpointer)
-    async def graph(state: dict) -> dict:
-        """React tool."""
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
 
-        values = []
+        @entrypoint(checkpointer=checkpointer)
+        async def graph(state: dict) -> dict:
+            """React tool."""
 
-        for idx in [1, 2, 3]:
-            values.extend([await double(idx), interrupt({"a": "boo"})])
+            values = []
 
-        return {"values": values}
+            for idx in [1, 2, 3]:
+                values.extend([await double(idx), interrupt({"a": "boo"})])
 
-    configurable = {"configurable": {"thread_id": uuid.uuid4()}}
-    await graph.ainvoke({}, configurable)
-    await graph.ainvoke(Command(resume="a"), configurable)
-    await graph.ainvoke(Command(resume="b"), configurable)
-    result = await graph.ainvoke(Command(resume="c"), configurable)
-    # `double` value should be cached appropriately when used w/ `interrupt`
-    assert result == {
-        "values": [2, "a", 4, "b", 6, "c"],
-    }
-    assert counter == 3
+            return {"values": values}
+
+        configurable = {"configurable": {"thread_id": uuid.uuid4()}}
+        await graph.ainvoke({}, configurable)
+        await graph.ainvoke(Command(resume="a"), configurable)
+        await graph.ainvoke(Command(resume="b"), configurable)
+        result = await graph.ainvoke(Command(resume="c"), configurable)
+        # `double` value should be cached appropriately when used w/ `interrupt`
+        assert result == {
+            "values": [2, "a", 4, "b", 6, "c"],
+        }
+        assert counter == 3