ci: Enable mypy checks for checkpoint-sqlite lib

This commit is contained in:
Nuno Campos
2024-09-19 08:46:44 -07:00
parent b8a8651c23
commit 6c0339ca4b
7 changed files with 72 additions and 53 deletions
+19 -15
View File
@@ -1,3 +1,5 @@
from typing import Any, cast
import pytest
from langchain_core.runnables import RunnableConfig
@@ -13,7 +15,7 @@ from langgraph.checkpoint.sqlite.utils import _metadata_predicate, search_where
class TestSqliteSaver:
@pytest.fixture(autouse=True)
def setup(self):
def setup(self) -> None:
# objects for test setup
self.config_1: RunnableConfig = {
"configurable": {
@@ -56,7 +58,7 @@ class TestSqliteSaver:
}
self.metadata_3: CheckpointMetadata = {}
def test_search(self):
def test_search(self) -> None:
with SqliteSaver.from_conn_string(":memory:") as saver:
# set up test
# save checkpoints
@@ -65,13 +67,13 @@ class TestSqliteSaver:
saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
# call method / assertions
query_1: CheckpointMetadata = {"source": "input"} # search by 1 key
query_2: CheckpointMetadata = {
query_1 = {"source": "input"} # search by 1 key
query_2 = {
"step": 1,
"writes": {"foo": "bar"},
} # search by multiple keys
query_3: CheckpointMetadata = {} # search by no keys, return all checkpoints
query_4: CheckpointMetadata = {"source": "update", "step": 1} # no match
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
query_4 = {"source": "update", "step": 1} # no match
search_results_1 = list(saver.list(None, filter=query_1))
assert len(search_results_1) == 1
@@ -99,16 +101,18 @@ class TestSqliteSaver:
# TODO: test before and limit params
def test_search_where(self):
def test_search_where(self) -> None:
# call method / assertions
expected_predicate_1 = "WHERE json_extract(CAST(metadata AS TEXT), '$.source') = ? AND json_extract(CAST(metadata AS TEXT), '$.step') = ? AND json_extract(CAST(metadata AS TEXT), '$.writes') = ? AND json_extract(CAST(metadata AS TEXT), '$.score') = ? AND checkpoint_id < ?"
expected_param_values_1 = ["input", 2, "{}", 1, "1"]
assert search_where(None, self.metadata_1, self.config_1) == (
assert search_where(
None, cast(dict[str, Any], self.metadata_1), self.config_1
) == (
expected_predicate_1,
expected_param_values_1,
)
def test_metadata_predicate(self):
def test_metadata_predicate(self) -> None:
# call method / assertions
expected_predicate_1 = [
"json_extract(CAST(metadata AS TEXT), '$.source') = ?",
@@ -122,26 +126,26 @@ class TestSqliteSaver:
"json_extract(CAST(metadata AS TEXT), '$.writes') = ?",
"json_extract(CAST(metadata AS TEXT), '$.score') IS ?",
]
expected_predicate_3 = []
expected_predicate_3: list[str] = []
expected_param_values_1 = ["input", 2, "{}", 1]
expected_param_values_2 = ["loop", 1, '{"foo":"bar"}', None]
expected_param_values_3 = []
expected_param_values_3: list[Any] = []
assert _metadata_predicate(self.metadata_1) == (
assert _metadata_predicate(cast(dict[str, Any], self.metadata_1)) == (
expected_predicate_1,
expected_param_values_1,
)
assert _metadata_predicate(self.metadata_2) == (
assert _metadata_predicate(cast(dict[str, Any], self.metadata_2)) == (
expected_predicate_2,
expected_param_values_2,
)
assert _metadata_predicate(self.metadata_3) == (
assert _metadata_predicate(cast(dict[str, Any], self.metadata_3)) == (
expected_predicate_3,
expected_param_values_3,
)
async def test_informative_async_errors(self):
async def test_informative_async_errors(self) -> None:
with SqliteSaver.from_conn_string(":memory:") as saver:
# call method / assertions
with pytest.raises(NotImplementedError, match="AsyncSqliteSaver"):