standard-tests[minor]: Add standard tests for BaseStore (#23360)

Add standard tests to base store abstraction. These only work on [str, str] right now. We'll need to check if it's possible to add encoder/decoders to generalize
1 week ago · 3b3ed72d35
parent e1190c8f3c
commit 3b3ed72d35
2 changed files with 306 additions and 0 deletions
--- a/libs/standard-tests/langchain_standard_tests/integration_tests/base_store.py
+++ b/libs/standard-tests/langchain_standard_tests/integration_tests/base_store.py
@ -0,0 +1,276 @@
+from abc import ABC, abstractmethod
+from typing import AsyncGenerator, Generator, Generic, Tuple, TypeVar
+
+import pytest
+from langchain_core.stores import BaseStore
+
+V = TypeVar("V")
+
+
+class BaseStoreSyncTests(ABC, Generic[V]):
+    """Test suite for checking the key-value API of a BaseStore.
+
+    This test suite verifies the basic key-value API of a BaseStore.
+
+    The test suite is designed for synchronous key-value stores.
+
+    Implementers should subclass this test suite and provide a fixture
+    that returns an empty key-value store for each test.
+    """
+
+    @abstractmethod
+    @pytest.fixture
+    def kv_store(self) -> BaseStore[str, V]:
+        """Get the key-value store class to test.
+
+        The returned key-value store should be EMPTY.
+        """
+
+    @abstractmethod
+    @pytest.fixture()
+    def three_values(self) -> Tuple[V, V, V]:
+        """Thee example values that will be used in the tests."""
+        pass
+
+    def test_three_values(self, three_values: Tuple[V, V, V]) -> None:
+        """Test that the fixture provides three values."""
+        assert isinstance(three_values, tuple)
+        assert len(three_values) == 3
+
+    def test_kv_store_is_empty(self, kv_store: BaseStore[str, V]) -> None:
+        """Test that the key-value store is empty."""
+        keys = ["foo", "bar", "buzz"]
+        assert kv_store.mget(keys) == [None, None, None]
+
+    def test_set_and_get_values(
+        self, kv_store: BaseStore[str, V], three_values: Tuple[V, V, V]
+    ) -> None:
+        """Test setting and getting values in the key-value store."""
+        foo = three_values[0]
+        bar = three_values[1]
+        key_value_pairs = [("foo", foo), ("bar", bar)]
+        kv_store.mset(key_value_pairs)
+        assert kv_store.mget(["foo", "bar"]) == [foo, bar]
+
+    def test_store_still_empty(self, kv_store: BaseStore[str, V]) -> None:
+        """This test should follow a test that sets values.
+
+        This just verifies that the fixture is set up properly to be empty
+        after each test.
+        """
+        keys = ["foo"]
+        assert kv_store.mget(keys) == [None]
+
+    def test_delete_values(
+        self, kv_store: BaseStore[str, V], three_values: Tuple[V, V, V]
+    ) -> None:
+        """Test deleting values from the key-value store."""
+        foo = three_values[0]
+        bar = three_values[1]
+        key_value_pairs = [("foo", foo), ("bar", bar)]
+        kv_store.mset(key_value_pairs)
+        kv_store.mdelete(["foo"])
+        assert kv_store.mget(["foo", "bar"]) == [None, bar]
+
+    def test_delete_bulk_values(
+        self, kv_store: BaseStore[str, V], three_values: Tuple[V, V, V]
+    ) -> None:
+        """Test that we can delete several values at once."""
+        foo, bar, buz = three_values
+        key_values = [("foo", foo), ("bar", bar), ("buz", buz)]
+        kv_store.mset(key_values)
+        kv_store.mdelete(["foo", "buz"])
+        assert kv_store.mget(["foo", "bar", "buz"]) == [None, bar, None]
+
+    def test_delete_missing_keys(self, kv_store: BaseStore[str, V]) -> None:
+        """Deleting missing keys should not raise an exception."""
+        kv_store.mdelete(["foo"])
+        kv_store.mdelete(["foo", "bar", "baz"])
+
+    def test_set_values_is_idempotent(
+        self, kv_store: BaseStore[str, V], three_values: Tuple[V, V, V]
+    ) -> None:
+        """Setting values by key should be idempotent."""
+        foo, bar, _ = three_values
+        key_value_pairs = [("foo", foo), ("bar", bar)]
+        kv_store.mset(key_value_pairs)
+        kv_store.mset(key_value_pairs)
+        assert kv_store.mget(["foo", "bar"]) == [foo, bar]
+        assert sorted(kv_store.yield_keys()) == ["bar", "foo"]
+
+    def test_get_can_get_same_value(
+        self, kv_store: BaseStore[str, V], three_values: Tuple[V, V, V]
+    ) -> None:
+        """Test that the same value can be retrieved multiple times."""
+        foo, bar, _ = three_values
+        key_value_pairs = [("foo", foo), ("bar", bar)]
+        kv_store.mset(key_value_pairs)
+        # This test assumes kv_store does not handle duplicates by default
+        assert kv_store.mget(["foo", "bar", "foo", "bar"]) == [foo, bar, foo, bar]
+
+    def test_overwrite_values_by_key(
+        self, kv_store: BaseStore[str, V], three_values: Tuple[V, V, V]
+    ) -> None:
+        """Test that we can overwrite values by key using mset."""
+        foo, bar, buzz = three_values
+        key_value_pairs = [("foo", foo), ("bar", bar)]
+        kv_store.mset(key_value_pairs)
+
+        # Now overwrite value of key "foo"
+        new_key_value_pairs = [("foo", buzz)]
+        kv_store.mset(new_key_value_pairs)
+
+        # Check that the value has been updated
+        assert kv_store.mget(["foo", "bar"]) == [buzz, bar]
+
+    def test_yield_keys(
+        self, kv_store: BaseStore[str, V], three_values: Tuple[V, V, V]
+    ) -> None:
+        """Test that we can yield keys from the store."""
+        foo, bar, buzz = three_values
+        key_value_pairs = [("foo", foo), ("bar", bar)]
+        kv_store.mset(key_value_pairs)
+
+        generator = kv_store.yield_keys()
+        assert isinstance(generator, Generator)
+
+        assert sorted(kv_store.yield_keys()) == ["bar", "foo"]
+        assert sorted(kv_store.yield_keys(prefix="foo")) == ["foo"]
+
+
+class BaseStoreAsyncTests(ABC):
+    """Test suite for checking the key-value API of a BaseStore.
+
+    This test suite verifies the basic key-value API of a BaseStore.
+
+    The test suite is designed for synchronous key-value stores.
+
+    Implementers should subclass this test suite and provide a fixture
+    that returns an empty key-value store for each test.
+    """
+
+    @abstractmethod
+    @pytest.fixture
+    async def kv_store(self) -> BaseStore[str, V]:
+        """Get the key-value store class to test.
+
+        The returned key-value store should be EMPTY.
+        """
+
+    @abstractmethod
+    @pytest.fixture()
+    def three_values(self) -> Tuple[V, V, V]:
+        """Thee example values that will be used in the tests."""
+        pass
+
+    async def test_three_values(self, three_values: Tuple[V, V, V]) -> None:
+        """Test that the fixture provides three values."""
+        assert isinstance(three_values, tuple)
+        assert len(three_values) == 3
+
+    async def test_kv_store_is_empty(self, kv_store: BaseStore[str, V]) -> None:
+        """Test that the key-value store is empty."""
+        keys = ["foo", "bar", "buzz"]
+        assert await kv_store.amget(keys) == [None, None, None]
+
+    async def test_set_and_get_values(
+        self, kv_store: BaseStore[str, V], three_values: Tuple[V, V, V]
+    ) -> None:
+        """Test setting and getting values in the key-value store."""
+        foo = three_values[0]
+        bar = three_values[1]
+        key_value_pairs = [("foo", foo), ("bar", bar)]
+        await kv_store.amset(key_value_pairs)
+        assert await kv_store.amget(["foo", "bar"]) == [foo, bar]
+
+    async def test_store_still_empty(self, kv_store: BaseStore[str, V]) -> None:
+        """This test should follow a test that sets values.
+
+        This just verifies that the fixture is set up properly to be empty
+        after each test.
+        """
+        keys = ["foo"]
+        assert await kv_store.amget(keys) == [None]
+
+    async def test_delete_values(
+        self, kv_store: BaseStore[str, V], three_values: Tuple[V, V, V]
+    ) -> None:
+        """Test deleting values from the key-value store."""
+        foo = three_values[0]
+        bar = three_values[1]
+        key_value_pairs = [("foo", foo), ("bar", bar)]
+        await kv_store.amset(key_value_pairs)
+        await kv_store.amdelete(["foo"])
+        assert await kv_store.amget(["foo", "bar"]) == [None, bar]
+
+    async def test_delete_bulk_values(
+        self, kv_store: BaseStore[str, V], three_values: Tuple[V, V, V]
+    ) -> None:
+        """Test that we can delete several values at once."""
+        foo, bar, buz = three_values
+        key_values = [("foo", foo), ("bar", bar), ("buz", buz)]
+        await kv_store.amset(key_values)
+        await kv_store.amdelete(["foo", "buz"])
+        assert await kv_store.amget(["foo", "bar", "buz"]) == [None, bar, None]
+
+    async def test_delete_missing_keys(self, kv_store: BaseStore[str, V]) -> None:
+        """Deleting missing keys should not raise an exception."""
+        await kv_store.amdelete(["foo"])
+        await kv_store.amdelete(["foo", "bar", "baz"])
+
+    async def test_set_values_is_idempotent(
+        self, kv_store: BaseStore[str, V], three_values: Tuple[V, V, V]
+    ) -> None:
+        """Setting values by key should be idempotent."""
+        foo, bar, _ = three_values
+        key_value_pairs = [("foo", foo), ("bar", bar)]
+        await kv_store.amset(key_value_pairs)
+        await kv_store.amset(key_value_pairs)
+        assert await kv_store.amget(["foo", "bar"]) == [foo, bar]
+        assert sorted(kv_store.yield_keys()) == ["bar", "foo"]
+
+    async def test_get_can_get_same_value(
+        self, kv_store: BaseStore[str, V], three_values: Tuple[V, V, V]
+    ) -> None:
+        """Test that the same value can be retrieved multiple times."""
+        foo, bar, _ = three_values
+        key_value_pairs = [("foo", foo), ("bar", bar)]
+        await kv_store.amset(key_value_pairs)
+        # This test assumes kv_store does not handle duplicates by async default
+        assert await kv_store.amget(["foo", "bar", "foo", "bar"]) == [
+            foo,
+            bar,
+            foo,
+            bar,
+        ]
+
+    async def test_overwrite_values_by_key(
+        self, kv_store: BaseStore[str, V], three_values: Tuple[V, V, V]
+    ) -> None:
+        """Test that we can overwrite values by key using mset."""
+        foo, bar, buzz = three_values
+        key_value_pairs = [("foo", foo), ("bar", bar)]
+        await kv_store.amset(key_value_pairs)
+
+        # Now overwrite value of key "foo"
+        new_key_value_pairs = [("foo", buzz)]
+        await kv_store.amset(new_key_value_pairs)
+
+        # Check that the value has been updated
+        assert await kv_store.amget(["foo", "bar"]) == [buzz, bar]
+
+    async def test_yield_keys(
+        self, kv_store: BaseStore[str, V], three_values: Tuple[V, V, V]
+    ) -> None:
+        """Test that we can yield keys from the store."""
+        foo, bar, buzz = three_values
+        key_value_pairs = [("foo", foo), ("bar", bar)]
+        await kv_store.amset(key_value_pairs)
+
+        generator = kv_store.ayield_keys()
+        assert isinstance(generator, AsyncGenerator)
+
+        assert sorted([key async for key in kv_store.ayield_keys()]) == ["bar", "foo"]
+        assert sorted([key async for key in kv_store.ayield_keys(prefix="foo")]) == [
+            "foo"
+        ]
--- a/libs/standard-tests/tests/unit_tests/test_in_memory_base_store.py
+++ b/libs/standard-tests/tests/unit_tests/test_in_memory_base_store.py
@ -0,0 +1,30 @@
+"""Tests for the InMemoryStore class."""
+from typing import Tuple
+
+import pytest
+from langchain_core.stores import InMemoryStore
+
+from langchain_standard_tests.integration_tests.base_store import (
+    BaseStoreAsyncTests,
+    BaseStoreSyncTests,
+)
+
+
+class TestInMemoryStore(BaseStoreSyncTests):
+    @pytest.fixture
+    def three_values(self) -> Tuple[str, str, str]:
+        return "foo", "bar", "buzz"
+
+    @pytest.fixture
+    def kv_store(self) -> InMemoryStore:
+        return InMemoryStore()
+
+
+class TestInMemoryStoreAsync(BaseStoreAsyncTests):
+    @pytest.fixture
+    def three_values(self) -> Tuple[str, str, str]:  # type: ignore
+        return "foo", "bar", "buzz"
+
+    @pytest.fixture
+    async def kv_store(self) -> InMemoryStore:
+        return InMemoryStore()