PR Cleanup

Author: strawgate

key-value/key-value-aio/src/key_value/aio/stores/elasticsearch/utils.py

@@ -121,7 +121,7 @@ class InstallSerializerMixin:
 
     @classmethod
     def install_serializer(cls, client: AsyncElasticsearch) -> None:
-        client._transport.serializers.serializers.update(
+        client.transport.serializers.serializers.update(
             {
                 cls.mimetype: cls(),
                 cls.compatibility_mimetype: cls(),
@@ -131,7 +131,7 @@ def install_serializer(cls, client: AsyncElasticsearch) -> None:
     @classmethod
     def install_default_serializer(cls, client: AsyncElasticsearch) -> None:
         cls.install_serializer(client=client)
-        client._transport.serializers.default_serializer = cls()
+        client.transport.serializers.default_serializer = cls()
 
 
 class LessCapableJsonSerializer(InstallSerializerMixin, JsonSerializer):

key-value/key-value-aio/tests/conftest.py

@@ -5,11 +5,11 @@
 import subprocess
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
-from time import sleep
 
 import pytest
 from docker import DockerClient
 from docker.models.containers import Container
+from key_value.shared.code_gen.sleep import sleep
 
 logger = logging.getLogger(__name__)
 

key-value/key-value-aio/tests/stores/elasticsearch/test_elasticsearch.py

@@ -145,7 +145,7 @@ async def test_put_put_two_indices(self, store: ElasticsearchStore, es_client: A
 
 @pytest.mark.skipif(should_skip_docker_tests(), reason="Docker is not running")
 class TestElasticsearchStoreNativeMode(BaseTestElasticsearchStore):
-    """Test Elasticsearch store in native mode"""
+    """Test Elasticsearch store in native mode (i.e. it stores flattened objects)"""
 
     @override
     @pytest.fixture
@@ -208,7 +208,7 @@ async def test_migration_from_legacy_mode(self, store: ElasticsearchStore, es_cl
 
 
 class TestElasticsearchStoreNonNativeMode(BaseTestElasticsearchStore):
-    """Test Elasticsearch store in non-native mode"""
+    """Test Elasticsearch store in non-native mode (i.e. it stores stringified JSON values)"""
 
     @override
     @pytest.fixture
@@ -220,8 +220,6 @@ async def test_value_stored_as_json_string(self, store: ElasticsearchStore, es_c
         """Verify values are stored as flattened objects, not JSON strings"""
         await store.put(collection="test", key="test_key", value={"name": "Alice", "age": 30})
 
-        # Check raw Elasticsearch document using public sanitization methods
-        # Note: We need to access these internal methods for testing the storage format
         index_name = store._sanitize_index_name(collection="test")  # pyright: ignore[reportPrivateUsage]
         doc_id = store._sanitize_document_id(key="test_key")  # pyright: ignore[reportPrivateUsage]
 
@@ -250,7 +248,6 @@ async def test_value_stored_as_json_string(self, store: ElasticsearchStore, es_c
 
     async def test_migration_from_native_mode(self, store: ElasticsearchStore, es_client: AsyncElasticsearch):
         """Verify non-native mode can read native mode data"""
-        # Manually insert a legacy document with JSON string value
         index_name = store._sanitize_index_name(collection="test")  # pyright: ignore[reportPrivateUsage]
         doc_id = store._sanitize_document_id(key="legacy_key")  # pyright: ignore[reportPrivateUsage]
 
@@ -266,6 +263,5 @@ async def test_migration_from_native_mode(self, store: ElasticsearchStore, es_cl
 
         await es_client.indices.refresh(index=index_name)
 
-        # Should be able to read it in native mode
         result = await store.get(collection="test", key="legacy_key")
         assert result == snapshot({"name": "Alice", "age": 30})

key-value/key-value-sync/src/key_value/sync/code_gen/stores/base.py

@@ -252,6 +252,7 @@ def _put_managed_entries(
         created_at: datetime,
         expires_at: datetime | None,
     ) -> None:
+
         """Store multiple managed entries by key in the specified collection.
 
         Args:

key-value/key-value-sync/src/key_value/sync/code_gen/stores/elasticsearch/utils.py

@@ -123,12 +123,12 @@ class InstallSerializerMixin:
 
     @classmethod
     def install_serializer(cls, client: Elasticsearch) -> None:
-        client._transport.serializers.serializers.update({cls.mimetype: cls(), cls.compatibility_mimetype: cls()})
+        client.transport.serializers.serializers.update({cls.mimetype: cls(), cls.compatibility_mimetype: cls()})
 
     @classmethod
     def install_default_serializer(cls, client: Elasticsearch) -> None:
         cls.install_serializer(client=client)
-        client._transport.serializers.default_serializer = cls()
+        client.transport.serializers.default_serializer = cls()
 
 
 class LessCapableJsonSerializer(InstallSerializerMixin, JsonSerializer):

key-value/key-value-sync/tests/code_gen/conftest.py

@@ -8,11 +8,11 @@
 import subprocess
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
-from time import sleep
 
 import pytest
 from docker import DockerClient
 from docker.models.containers import Container
+from key_value.shared.code_gen.sleep import sleep
 
 logger = logging.getLogger(__name__)
 

key-value/key-value-sync/tests/code_gen/stores/elasticsearch/test_elasticsearch.py

@@ -146,7 +146,7 @@ def test_put_put_two_indices(self, store: ElasticsearchStore, es_client: Elastic
 
 @pytest.mark.skipif(should_skip_docker_tests(), reason="Docker is not running")
 class TestElasticsearchStoreNativeMode(BaseTestElasticsearchStore):
-    """Test Elasticsearch store in native mode"""
+    """Test Elasticsearch store in native mode (i.e. it stores flattened objects)"""
 
     @override
     @pytest.fixture
@@ -202,7 +202,7 @@ def test_migration_from_legacy_mode(self, store: ElasticsearchStore, es_client:
 
 
 class TestElasticsearchStoreNonNativeMode(BaseTestElasticsearchStore):
-    """Test Elasticsearch store in non-native mode"""
+    """Test Elasticsearch store in non-native mode (i.e. it stores stringified JSON values)"""
 
     @override
     @pytest.fixture
@@ -214,8 +214,6 @@ def test_value_stored_as_json_string(self, store: ElasticsearchStore, es_client:
         """Verify values are stored as flattened objects, not JSON strings"""
         store.put(collection="test", key="test_key", value={"name": "Alice", "age": 30})
 
-        # Check raw Elasticsearch document using public sanitization methods
-        # Note: We need to access these internal methods for testing the storage format
         index_name = store._sanitize_index_name(collection="test")  # pyright: ignore[reportPrivateUsage]
         doc_id = store._sanitize_document_id(key="test_key")  # pyright: ignore[reportPrivateUsage]
 
@@ -244,7 +242,6 @@ def test_value_stored_as_json_string(self, store: ElasticsearchStore, es_client:
 
     def test_migration_from_native_mode(self, store: ElasticsearchStore, es_client: Elasticsearch):
         """Verify non-native mode can read native mode data"""
-        # Manually insert a legacy document with JSON string value
         index_name = store._sanitize_index_name(collection="test")  # pyright: ignore[reportPrivateUsage]
         doc_id = store._sanitize_document_id(key="legacy_key")  # pyright: ignore[reportPrivateUsage]
 
@@ -256,6 +253,5 @@ def test_migration_from_native_mode(self, store: ElasticsearchStore, es_client:
 
         es_client.indices.refresh(index=index_name)
 
-        # Should be able to read it in native mode
         result = store.get(collection="test", key="legacy_key")
         assert result == snapshot({"name": "Alice", "age": 30})