Fix type issues in tests

This commit is contained in:
Nuno Campos
2024-09-18 18:16:51 -07:00
parent 4d26c3599f
commit 333c5f5501
+13 -11
View File
@@ -1,3 +1,5 @@
from typing import Any
import pytest
from langchain_core.runnables import RunnableConfig
@@ -12,7 +14,7 @@ from langgraph.checkpoint.memory import MemorySaver
class TestMemorySaver:
@pytest.fixture(autouse=True)
def setup(self):
def setup(self) -> None:
self.memory_saver = MemorySaver()
# objects for test setup
@@ -57,7 +59,7 @@ class TestMemorySaver:
}
self.metadata_3: CheckpointMetadata = {}
async def test_search(self):
async def test_search(self) -> None:
# set up test
# save checkpoints
self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
@@ -65,13 +67,13 @@ class TestMemorySaver:
self.memory_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(self.memory_saver.list(None, filter=query_1))
assert len(search_results_1) == 1
@@ -99,7 +101,7 @@ class TestMemorySaver:
# TODO: test before and limit params
async def test_asearch(self):
async def test_asearch(self) -> None:
# set up test
# save checkpoints
self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
@@ -107,13 +109,13 @@ class TestMemorySaver:
self.memory_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 = [
c async for c in self.memory_saver.alist(None, filter=query_1)