updates for opensearch tests

Author: strawgate

key-value/key-value-aio/src/key_value/aio/stores/opensearch/store.py

@@ -68,8 +68,8 @@
         },
         "value": {
             "properties": {
-                "flattened": {
-                    "type": "flattened",
+                "flat": {
+                    "type": "flat_object",
                 },
             },
         },
@@ -101,14 +101,14 @@ def prepare_dump(self, data: dict[str, Any]) -> dict[str, Any]:
         value = data.pop("value")
 
         data["value"] = {
-            "flattened": value,
+            "flat": value,
         }
 
         return data
 
     @override
     def prepare_load(self, data: dict[str, Any]) -> dict[str, Any]:
-        data["value"] = data.pop("value").get("flattened")
+        data["value"] = data.pop("value").get("flat")
 
         return data
 
@@ -224,7 +224,6 @@ def __init__(
         if opensearch_client:
             self._client = opensearch_client
         elif url:
-            # Build kwargs for AsyncOpenSearch
             client_kwargs: dict[str, Any] = {
                 "hosts": [url],
                 "http_compress": True,

key-value/key-value-aio/tests/stores/opensearch/test_opensearch.py

@@ -1,15 +1,18 @@
 import contextlib
 from collections.abc import AsyncGenerator
 from datetime import datetime, timedelta, timezone
+from typing import Any
 
 import pytest
-from dirty_equals import IsFloat
+from dirty_equals import IsFloat, IsStr
+from elasticsearch import AsyncElasticsearch
 from inline_snapshot import snapshot
 from key_value.shared.stores.wait import async_wait_for_true
 from key_value.shared.utils.managed_entry import ManagedEntry
 from opensearchpy import AsyncOpenSearch
 from typing_extensions import override
 
+from key_value.aio.protocols.key_value import AsyncKeyValueProtocol
 from key_value.aio.stores.base import BaseStore
 from key_value.aio.stores.opensearch import OpenSearchStore
 from key_value.aio.stores.opensearch.store import (
@@ -21,10 +24,13 @@
 from tests.stores.base import BaseStoreTests, ContextManagerStoreTestMixin
 
 TEST_SIZE_LIMIT = 1 * 1024 * 1024  # 1MB
-OS_HOST = "localhost"
-OS_PORT = 9200
-OS_URL = f"http://{OS_HOST}:{OS_PORT}"
-OS_CONTAINER_PORT = 9200
+LOCALHOST = "localhost"
+
+CONTAINER_PORT = 9200
+HOST_PORT = 19200
+
+OPENSEARCH_URL = f"http://{LOCALHOST}:{HOST_PORT}"
+
 
 WAIT_FOR_OPENSEARCH_TIMEOUT = 30
 
@@ -35,15 +41,15 @@
 
 
 def get_opensearch_client() -> AsyncOpenSearch:
-    return AsyncOpenSearch(hosts=[OS_URL], use_ssl=False, verify_certs=False)
+    return AsyncOpenSearch(hosts=[OPENSEARCH_URL], use_ssl=False, verify_certs=False)
 
 
 async def ping_opensearch() -> bool:
-    os_client: AsyncOpenSearch = get_opensearch_client()
+    opensearch_client: AsyncOpenSearch = get_opensearch_client()
 
-    async with os_client:
+    async with opensearch_client:
         try:
-            return await os_client.ping()
+            return await opensearch_client.ping()
         except Exception:
             return False
 
@@ -69,7 +75,7 @@ def test_managed_entry_document_conversion():
 
     assert document == snapshot(
         {
-            "value": {"flattened": {"test": "test"}},
+            "value": {"f": {"test": "test"}},
             "created_at": "2025-01-01T00:00:00+00:00",
             "expires_at": "2025-01-01T00:00:10+00:00",
         }
@@ -94,7 +100,7 @@ async def setup_opensearch(self, request: pytest.FixtureRequest) -> AsyncGenerat
         with docker_container(
             f"opensearch-test-{version}",
             os_image,
-            {str(OS_CONTAINER_PORT): OS_PORT},
+            {str(CONTAINER_PORT): HOST_PORT},
             {
                 "discovery.type": "single-node",
                 "DISABLE_SECURITY_PLUGIN": "true",
@@ -109,16 +115,16 @@ async def setup_opensearch(self, request: pytest.FixtureRequest) -> AsyncGenerat
 
     @pytest.fixture
     async def opensearch_client(self, setup_opensearch: None) -> AsyncGenerator[AsyncOpenSearch, None]:
-        os_client = get_opensearch_client()
+        opensearch_client = get_opensearch_client()
 
-        async with os_client:
-            await cleanup_opensearch_indices(opensearch_client=os_client)
+        async with opensearch_client:
+            await cleanup_opensearch_indices(opensearch_client=opensearch_client)
 
-            yield os_client
+            yield opensearch_client
 
     @override
     @pytest.fixture
-    async def default_store(self, opensearch_client: AsyncOpenSearch) -> AsyncGenerator[BaseStore, None]:
+    async def store(self, opensearch_client: AsyncOpenSearch) -> AsyncGenerator[BaseStore, None]:
         store = OpenSearchStore(
             opensearch_client=opensearch_client,
             index_prefix="opensearch-kv-store-e2e-test",
@@ -130,40 +136,86 @@ async def default_store(self, opensearch_client: AsyncOpenSearch) -> AsyncGenera
 
     @override
     @pytest.fixture
-    async def collection_sanitized_store(self, opensearch_client: AsyncOpenSearch) -> AsyncGenerator[BaseStore, None]:
+    async def sanitizing_store(self, opensearch_client: AsyncOpenSearch) -> AsyncGenerator[BaseStore, None]:
         store = OpenSearchStore(
             opensearch_client=opensearch_client,
             index_prefix="opensearch-kv-store-e2e-test",
             default_collection="test-collection",
+            key_sanitization_strategy=OpenSearchV1KeySanitizationStrategy(),
             collection_sanitization_strategy=OpenSearchV1CollectionSanitizationStrategy(),
         )
 
         async with store:
             yield store
 
     @override
-    @pytest.fixture
-    async def key_sanitized_store(self, opensearch_client: AsyncOpenSearch) -> AsyncGenerator[BaseStore, None]:
-        store = OpenSearchStore(
-            opensearch_client=opensearch_client,
-            index_prefix="opensearch-kv-store-e2e-test",
-            default_collection="test-collection",
-            key_sanitization_strategy=OpenSearchV1KeySanitizationStrategy(),
-        )
+    @pytest.mark.timeout(120)
+    async def test_store(self, store: BaseStore):
+        """Tests that the store is a valid AsyncKeyValueProtocol."""
+        assert isinstance(store, AsyncKeyValueProtocol) is True
 
-        async with store:
-            yield store
+    @pytest.mark.skip(reason="Distributed Caches are unbounded")
+    @override
+    async def test_not_unbounded(self, store: BaseStore): ...
 
+    @pytest.mark.skip(reason="Skip concurrent tests on distributed caches")
     @override
-    @pytest.fixture
-    async def fully_sanitized_store(self, opensearch_client: AsyncOpenSearch) -> AsyncGenerator[BaseStore, None]:
-        store = OpenSearchStore(
-            opensearch_client=opensearch_client,
-            index_prefix="opensearch-kv-store-e2e-test",
-            default_collection="test-collection",
-            key_sanitization_strategy=OpenSearchV1KeySanitizationStrategy(),
-            collection_sanitization_strategy=OpenSearchV1CollectionSanitizationStrategy(),
+    async def test_concurrent_operations(self, store: BaseStore): ...
+
+    @override
+    async def test_long_collection_name(self, store: OpenSearchStore, sanitizing_store: OpenSearchStore):  # pyright: ignore[reportIncompatibleMethodOverride]
+        with pytest.raises(Exception):  # noqa: B017, PT011
+            await store.put(collection="test_collection" * 100, key="test_key", value={"test": "test"})
+
+        await sanitizing_store.put(collection="test_collection" * 100, key="test_key", value={"test": "test"})
+        assert await sanitizing_store.get(collection="test_collection" * 100, key="test_key") == {"test": "test"}
+
+    @override
+    async def test_long_key_name(self, store: OpenSearchStore, sanitizing_store: OpenSearchStore):  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Tests that a long key name will not raise an error."""
+        with pytest.raises(Exception):  # noqa: B017, PT011
+            await store.put(collection="test_collection", key="test_key" * 100, value={"test": "test"})
+
+        await sanitizing_store.put(collection="test_collection", key="test_key" * 100, value={"test": "test"})
+        assert await sanitizing_store.get(collection="test_collection", key="test_key" * 100) == {"test": "test"}
+
+    async def test_put_put_two_indices(self, store: OpenSearchStore, opensearch_client: AsyncOpenSearch):
+        await store.put(collection="test_collection", key="test_key", value={"test": "test"})
+        await store.put(collection="test_collection_2", key="test_key", value={"test": "test"})
+        assert await store.get(collection="test_collection", key="test_key") == {"test": "test"}
+        assert await store.get(collection="test_collection_2", key="test_key") == {"test": "test"}
+
+        indices: dict[str, Any] = await opensearch_client.indices.get(index="opensearch-kv-store-e2e-test-*")
+        index_names: list[str] = list(indices.keys())
+        assert index_names == snapshot(["opensearch-kv-store-e2e-test-test_collection", "opensearch-kv-store-e2e-test-test_collection_2"])
+
+    async def test_value_stored_as_f_object(self, store: OpenSearchStore, opensearch_client: AsyncElasticsearch):
+        """Verify values are stored as f objects, not JSON strings"""
+        await store.put(collection="test", key="test_key", value={"name": "Alice", "age": 30})
+
+        index_name = store._get_index_name(collection="test")  # pyright: ignore[reportPrivateUsage]
+        doc_id = store._get_document_id(key="test_key")  # pyright: ignore[reportPrivateUsage]
+
+        response = await opensearch_client.get(index=index_name, id=doc_id)
+        assert response.body["_source"] == snapshot(
+            {
+                "value": {"f": {"name": "Alice", "age": 30}},
+                "created_at": IsStr(min_length=20, max_length=40),
+            }
         )
 
-        async with store:
-            yield store
+        # Test with TTL
+        await store.put(collection="test", key="test_key", value={"name": "Bob", "age": 25}, ttl=10)
+        response = await opensearch_client.get(index=index_name, id=doc_id)
+        assert response.body["_source"] == snapshot(
+            {
+                "value": {"f": {"name": "Bob", "age": 25}},
+                "created_at": IsStr(min_length=20, max_length=40),
+                "expires_at": IsStr(min_length=20, max_length=40),
+            }
+        )
+
+    @override
+    async def test_special_characters_in_collection_name(self, store: OpenSearchStore, sanitizing_store: OpenSearchStore):  # pyright: ignore[reportIncompatibleMethodOverride]
+        """Tests that a special characters in the collection name will not raise an error."""
+        await super().test_special_characters_in_collection_name(store=sanitizing_store)

key-value/key-value-sync/src/key_value/sync/code_gen/stores/opensearch/store.py

@@ -50,7 +50,7 @@
         "expires_at": {"type": "date"},
         "collection": {"type": "keyword"},
         "key": {"type": "keyword"},
-        "value": {"properties": {"flattened": {"type": "flattened"}}},
+        "value": {"properties": {"flat": {"type": "flat_object"}}},
     }
 }
 
@@ -78,13 +78,13 @@ def __init__(self) -> None:
     def prepare_dump(self, data: dict[str, Any]) -> dict[str, Any]:
         value = data.pop("value")
 
-        data["value"] = {"flattened": value}
+        data["value"] = {"flat": value}
 
         return data
 
     @override
     def prepare_load(self, data: dict[str, Any]) -> dict[str, Any]:
-        data["value"] = data.pop("value").get("flattened")
+        data["value"] = data.pop("value").get("flat")
 
         return data
 
@@ -198,7 +198,6 @@ def __init__(
         if opensearch_client:
             self._client = opensearch_client
         elif url:
-            # Build kwargs for AsyncOpenSearch
             client_kwargs: dict[str, Any] = {"hosts": [url], "http_compress": True, "timeout": 10, "max_retries": 3}
             if api_key:
                 client_kwargs["api_key"] = api_key