test: Ensure that async checkpointers in tests are setup in same event loop as they are used in

- This more closely remembers the environment they're used in, so it's what we should be testing
- Remove unnecessary pytest-asyncio dependency, use anyio pytest plugin instead
- Convert remaining async tests using only memory checkpointer to use all existing ones
This commit is contained in:
Nuno Campos
2024-09-02 17:24:54 -07:00
parent a937754132
commit d65d0ec665
14 changed files with 5975 additions and 4414 deletions
+21 -20
View File
@@ -4,12 +4,16 @@ import pytest
from pytest_mock import MockerFixture
from langgraph.graph import END, START, StateGraph
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite", "postgres", "postgres_pipe"],
from tests.conftest import (
ALL_CHECKPOINTERS_ASYNC,
ALL_CHECKPOINTERS_SYNC,
awith_checkpointer,
)
pytestmark = pytest.mark.anyio
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_interruption_without_state_updates(
request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture
) -> None:
@@ -47,12 +51,9 @@ def test_interruption_without_state_updates(
assert graph.get_state(thread).next == ()
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"],
)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_interruption_without_state_updates_async(
request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture
checkpointer_name: str, mocker: MockerFixture
):
"""Test interruption without state updates. This test confirms that
interrupting doesn't require a state key having been updated in the prev step"""
@@ -72,17 +73,17 @@ async def test_interruption_without_state_updates_async(
builder.add_edge("step_2", "step_3")
builder.add_edge("step_3", END)
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
graph = builder.compile(checkpointer=checkpointer, interrupt_after="*")
async with awith_checkpointer(checkpointer_name) as checkpointer:
graph = builder.compile(checkpointer=checkpointer, interrupt_after="*")
initial_input = {"input": "hello world"}
thread = {"configurable": {"thread_id": "1"}}
initial_input = {"input": "hello world"}
thread = {"configurable": {"thread_id": "1"}}
await graph.ainvoke(initial_input, thread, debug=True)
assert (await graph.aget_state(thread)).next == ("step_2",)
await graph.ainvoke(initial_input, thread, debug=True)
assert (await graph.aget_state(thread)).next == ("step_2",)
await graph.ainvoke(None, thread, debug=True)
assert (await graph.aget_state(thread)).next == ("step_3",)
await graph.ainvoke(None, thread, debug=True)
assert (await graph.aget_state(thread)).next == ("step_3",)
await graph.ainvoke(None, thread, debug=True)
assert (await graph.aget_state(thread)).next == ()
await graph.ainvoke(None, thread, debug=True)
assert (await graph.aget_state(thread)).next == ()