merge

libs/checkpoint-postgres/tests/conftest.py

@@ -1,4 +1,4 @@
-from typing import AsyncIterator
+from collections.abc import AsyncIterator
 
 import pytest
 from psycopg import AsyncConnection

libs/checkpoint-postgres/tests/test_async_store.py

@@ -1,7 +1,7 @@
 # type: ignore
 import sys
 import uuid
-from typing import AsyncIterator
+from collections.abc import AsyncIterator
 
 import pytest
 from conftest import DEFAULT_URI  # type: ignore
@@ -43,9 +43,8 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]:
                 yield store
         elif request.param == "pool":
             async with AsyncPostgresStore.from_conn_string(
-                conn_string, use_pool=True, max_size=10
+                conn_string, pool_config={"min_size": 1, "max_size": 10}
             ) as store:
-                await store.setup()
                 yield store
         else:  # default
             async with AsyncPostgresStore.from_conn_string(conn_string) as store:

libs/checkpoint-postgres/tests/test_store.py

@@ -7,7 +7,6 @@ import pytest
 from conftest import DEFAULT_URI  # type: ignore
 from langchain_core.embeddings import Embeddings
 from psycopg import Connection
-from psycopg_pool import ConnectionPool
 from utils import CharacterEmbeddings
 
 from langgraph.store.base import (
@@ -45,10 +44,9 @@ def store(request) -> PostgresStore:
             with PostgresStore.from_conn_string(conn_string, pipeline=True) as store:
                 yield store
         elif request.param == "pool":
-            with ConnectionPool(
-                conn_string, max_size=10, kwargs={"autocommit": True}
-            ) as pool:
-                store = PostgresStore(pool)
+            with PostgresStore.from_conn_string(
+                conn_string, pool_config={"min_size": 1, "max_size": 10}
+            ) as store:
                 yield store
         else:  # default
             with PostgresStore.from_conn_string(conn_string) as store:
@@ -566,7 +564,7 @@ def test_extract_text_by_path():
     assert _extract_text_by_path(nested_data, "empty_dict") == ["{}"]
 
     zeros = _extract_text_by_path(nested_data, "zeros[*]")
-    assert set(zeros) == {"0", "0.0", "0"}
+    assert set(zeros) == {"0", "0.0"}
 
     assert _extract_text_by_path(nested_data, "items[].value") == []
     assert _extract_text_by_path(nested_data, "items[abc].value") == []