Updates for tests and error handling

Author: strawgate

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

@@ -354,8 +354,13 @@ async def _put_managed_entries(
             )
 
             operations.extend([index_action, document])
-
-        _ = await self._client.bulk(operations=operations, refresh=self._should_refresh_on_put)  # pyright: ignore[reportUnknownMemberType]
+        try:
+            _ = await self._client.bulk(operations=operations, refresh=self._should_refresh_on_put)  # pyright: ignore[reportUnknownMemberType]
+        except ElasticsearchSerializationError as e:
+            msg = f"Failed to serialize bulk operations: {e}"
+            raise SerializationError(message=msg) from e
+        except Exception:
+            raise
 
     @override
     async def _delete_managed_entry(self, *, key: str, collection: str) -> bool:

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

@@ -37,7 +37,7 @@ def get_aggregations_from_body(body: dict[str, Any]) -> dict[str, Any]:
     if not isinstance(aggregations, dict) or not all(
         isinstance(key, str)
         for key in aggregations  # pyright: ignore[reportUnknownVariableType]
-    ):  # pyright: ignore[reportUnknownVariableType]
+    ):
         return {}
 
     return cast("dict[str, Any]", aggregations)

key-value/key-value-aio/tests/stores/redis/test_redis.py

@@ -36,7 +36,7 @@ def test_managed_entry_document_conversion():
     document = managed_entry_to_json(managed_entry=managed_entry)
 
     assert document == snapshot(
-        '{"value": {"test": "test"}, "created_at": "2025-01-01T00:00:00+00:00", "expires_at": "2025-01-01T00:00:10+00:00"}'
+        '{"created_at": "2025-01-01T00:00:00+00:00", "expires_at": "2025-01-01T00:00:10+00:00", "value": {"test": "test"}}'
     )
 
     round_trip_managed_entry = json_to_managed_entry(json_str=document)

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

@@ -303,8 +303,13 @@ def _put_managed_entries(
             )
 
             operations.extend([index_action, document])
-
-        _ = self._client.bulk(operations=operations, refresh=self._should_refresh_on_put)  # pyright: ignore[reportUnknownMemberType]
+        try:
+            _ = self._client.bulk(operations=operations, refresh=self._should_refresh_on_put)  # pyright: ignore[reportUnknownMemberType]
+        except ElasticsearchSerializationError as e:
+            msg = f"Failed to serialize bulk operations: {e}"
+            raise SerializationError(message=msg) from e
+        except Exception:
+            raise
 
     @override
     def _delete_managed_entry(self, *, key: str, collection: str) -> bool:

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

@@ -32,7 +32,6 @@ def get_aggregations_from_body(body: dict[str, Any]) -> dict[str, Any]:
         return {}
 
     if not isinstance(aggregations, dict) or not all(isinstance(key, str) for key in aggregations):  # pyright: ignore[reportUnknownVariableType]
-        # pyright: ignore[reportUnknownVariableType]
         return {}
 
     return cast("dict[str, Any]", aggregations)

key-value/key-value-sync/tests/code_gen/stores/redis/test_redis.py

@@ -36,7 +36,7 @@ def test_managed_entry_document_conversion():
     document = managed_entry_to_json(managed_entry=managed_entry)
 
     assert document == snapshot(
-        '{"value": {"test": "test"}, "created_at": "2025-01-01T00:00:00+00:00", "expires_at": "2025-01-01T00:00:10+00:00"}'
+        '{"created_at": "2025-01-01T00:00:00+00:00", "expires_at": "2025-01-01T00:00:10+00:00", "value": {"test": "test"}}'
     )
 
     round_trip_managed_entry = json_to_managed_entry(json_str=document)