Skip to content

Commit d823f90

Browse files
strawgatestrawgate
committed
ES Test refactor
1 parent 65b755d commit d823f90

File tree

4 files changed

+57
-50
lines changed

4 files changed

+57
-50
lines changed

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

Lines changed: 28 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from datetime import datetime, timedelta, timezone
33

44
import pytest
5-
from dirty_equals import IsFloat
5+
from dirty_equals import IsFloat, IsStr
66
from elasticsearch import AsyncElasticsearch
77
from inline_snapshot import snapshot
88
from key_value.shared.stores.wait import async_wait_for_true
@@ -94,8 +94,7 @@ def test_managed_entry_document_conversion_native_storage():
9494
assert round_trip_managed_entry.expires_at == expires_at
9595

9696

97-
@pytest.mark.skipif(should_skip_docker_tests(), reason="Docker is not running")
98-
class TestElasticsearchStore(ContextManagerStoreTestMixin, BaseStoreTests):
97+
class BaseTestElasticsearchStore(ContextManagerStoreTestMixin, BaseStoreTests):
9998
@pytest.fixture(autouse=True, scope="session", params=ELASTICSEARCH_VERSIONS_TO_TEST)
10099
async def setup_elasticsearch(self, request: pytest.FixtureRequest) -> AsyncGenerator[None, None]:
101100
version = request.param
@@ -124,13 +123,6 @@ async def _cleanup(self):
124123
for index in indices:
125124
_ = await elasticsearch_client.options(ignore_status=404).indices.delete(index=index)
126125

127-
@override
128-
@pytest.fixture
129-
async def store(self) -> AsyncGenerator[ElasticsearchStore, None]:
130-
await self._cleanup()
131-
async with ElasticsearchStore(url=ES_URL, index_prefix="kv-store-e2e-test", native_storage=True) as store:
132-
yield store
133-
134126
@pytest.mark.skip(reason="Distributed Caches are unbounded")
135127
@override
136128
async def test_not_unbounded(self, store: BaseStore): ...
@@ -150,6 +142,17 @@ async def test_put_put_two_indices(self, store: ElasticsearchStore, es_client: A
150142
assert "kv-store-e2e-test-test_collection" in indices
151143
assert "kv-store-e2e-test-test_collection_2" in indices
152144

145+
146+
@pytest.mark.skipif(should_skip_docker_tests(), reason="Docker is not running")
147+
class TestElasticsearchStoreNativeMode(BaseTestElasticsearchStore):
148+
"""Test Elasticsearch store in native mode"""
149+
150+
@override
151+
@pytest.fixture
152+
async def store(self) -> ElasticsearchStore:
153+
await self._cleanup()
154+
return ElasticsearchStore(url=ES_URL, index_prefix="kv-store-e2e-test", native_storage=True)
155+
153156
async def test_value_stored_as_flattened_object(self, store: ElasticsearchStore, es_client: AsyncElasticsearch):
154157
"""Verify values are stored as flattened objects, not JSON strings"""
155158
await store.put(collection="test", key="test_key", value={"name": "Alice", "age": 30})
@@ -164,8 +167,8 @@ async def test_value_stored_as_flattened_object(self, store: ElasticsearchStore,
164167
{
165168
"collection": "test",
166169
"key": "test_key",
167-
"value": {"string": '{"name": "Alice", "age": 30}'},
168-
"created_at": "2025-10-27T15:33:31.795253+00:00",
170+
"value": {"flattened": {"name": "Alice", "age": 30}},
171+
"created_at": IsStr(min_length=20, max_length=40),
169172
}
170173
)
171174

@@ -176,9 +179,9 @@ async def test_value_stored_as_flattened_object(self, store: ElasticsearchStore,
176179
{
177180
"collection": "test",
178181
"key": "test_key",
179-
"value": {"string": '{"name": "Bob", "age": 25}'},
180-
"created_at": "2025-10-27T15:33:32.040020+00:00",
181-
"expires_at": "2025-10-27T15:33:42.040020+00:00",
182+
"value": {"flattened": {"name": "Bob", "age": 25}},
183+
"created_at": IsStr(min_length=20, max_length=40),
184+
"expires_at": IsStr(min_length=20, max_length=40),
182185
}
183186
)
184187

@@ -204,13 +207,14 @@ async def test_migration_from_legacy_mode(self, store: ElasticsearchStore, es_cl
204207
assert result == snapshot(None)
205208

206209

207-
class TestElasticsearchStoreNonNativeMode(TestElasticsearchStore):
210+
class TestElasticsearchStoreNonNativeMode(BaseTestElasticsearchStore):
211+
"""Test Elasticsearch store in non-native mode"""
212+
208213
@override
209214
@pytest.fixture
210-
async def store(self) -> AsyncGenerator[ElasticsearchStore, None]:
215+
async def store(self) -> ElasticsearchStore:
211216
await self._cleanup()
212-
async with ElasticsearchStore(url=ES_URL, index_prefix="kv-store-e2e-test", native_storage=False) as store:
213-
yield store
217+
return ElasticsearchStore(url=ES_URL, index_prefix="kv-store-e2e-test", native_storage=False)
214218

215219
async def test_value_stored_as_json_string(self, store: ElasticsearchStore, es_client: AsyncElasticsearch):
216220
"""Verify values are stored as flattened objects, not JSON strings"""
@@ -226,8 +230,8 @@ async def test_value_stored_as_json_string(self, store: ElasticsearchStore, es_c
226230
{
227231
"collection": "test",
228232
"key": "test_key",
229-
"value": {"string": '{"name": "Alice", "age": 30}'},
230-
"created_at": "2025-10-27T15:33:32.631007+00:00",
233+
"value": {"string": '{"age": 30, "name": "Alice"}'},
234+
"created_at": IsStr(min_length=20, max_length=40),
231235
}
232236
)
233237

@@ -238,9 +242,9 @@ async def test_value_stored_as_json_string(self, store: ElasticsearchStore, es_c
238242
{
239243
"collection": "test",
240244
"key": "test_key",
241-
"value": {"string": '{"name": "Bob", "age": 25}'},
242-
"created_at": "2025-10-27T15:33:32.650061+00:00",
243-
"expires_at": "2025-10-27T15:33:42.650061+00:00",
245+
"value": {"string": '{"age": 25, "name": "Bob"}'},
246+
"created_at": IsStr(min_length=20, max_length=40),
247+
"expires_at": IsStr(min_length=20, max_length=40),
244248
}
245249
)
246250

‎key-value/key-value-shared/src/key_value/shared/utils/managed_entry.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,7 @@ def from_json(cls, json_str: str, includes_metadata: bool = True, ttl: SupportsF
125125

126126
def dump_to_json(obj: dict[str, Any]) -> str:
127127
try:
128-
return json.dumps(obj)
128+
return json.dumps(obj, sort_keys=True)
129129
except (json.JSONDecodeError, TypeError) as e:
130130
msg: str = f"Failed to serialize object to JSON: {e}"
131131
raise SerializationError(msg) from e

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -252,7 +252,6 @@ def _put_managed_entries(
252252
created_at: datetime,
253253
expires_at: datetime | None,
254254
) -> None:
255-
256255
"""Store multiple managed entries by key in the specified collection.
257256
258257
Args:

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

Lines changed: 28 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from datetime import datetime, timedelta, timezone
66

77
import pytest
8-
from dirty_equals import IsFloat
8+
from dirty_equals import IsFloat, IsStr
99
from elasticsearch import Elasticsearch
1010
from inline_snapshot import snapshot
1111
from key_value.shared.stores.wait import wait_for_true
@@ -95,8 +95,7 @@ def test_managed_entry_document_conversion_native_storage():
9595
assert round_trip_managed_entry.expires_at == expires_at
9696

9797

98-
@pytest.mark.skipif(should_skip_docker_tests(), reason="Docker is not running")
99-
class TestElasticsearchStore(ContextManagerStoreTestMixin, BaseStoreTests):
98+
class BaseTestElasticsearchStore(ContextManagerStoreTestMixin, BaseStoreTests):
10099
@pytest.fixture(autouse=True, scope="session", params=ELASTICSEARCH_VERSIONS_TO_TEST)
101100
def setup_elasticsearch(self, request: pytest.FixtureRequest) -> Generator[None, None, None]:
102101
version = request.param
@@ -125,13 +124,6 @@ def _cleanup(self):
125124
for index in indices:
126125
_ = elasticsearch_client.options(ignore_status=404).indices.delete(index=index)
127126

128-
@override
129-
@pytest.fixture
130-
def store(self) -> Generator[ElasticsearchStore, None, None]:
131-
self._cleanup()
132-
with ElasticsearchStore(url=ES_URL, index_prefix="kv-store-e2e-test", native_storage=True) as store:
133-
yield store
134-
135127
@pytest.mark.skip(reason="Distributed Caches are unbounded")
136128
@override
137129
def test_not_unbounded(self, store: BaseStore): ...
@@ -151,6 +143,17 @@ def test_put_put_two_indices(self, store: ElasticsearchStore, es_client: Elastic
151143
assert "kv-store-e2e-test-test_collection" in indices
152144
assert "kv-store-e2e-test-test_collection_2" in indices
153145

146+
147+
@pytest.mark.skipif(should_skip_docker_tests(), reason="Docker is not running")
148+
class TestElasticsearchStoreNativeMode(BaseTestElasticsearchStore):
149+
"""Test Elasticsearch store in native mode"""
150+
151+
@override
152+
@pytest.fixture
153+
def store(self) -> ElasticsearchStore:
154+
self._cleanup()
155+
return ElasticsearchStore(url=ES_URL, index_prefix="kv-store-e2e-test", native_storage=True)
156+
154157
def test_value_stored_as_flattened_object(self, store: ElasticsearchStore, es_client: Elasticsearch):
155158
"""Verify values are stored as flattened objects, not JSON strings"""
156159
store.put(collection="test", key="test_key", value={"name": "Alice", "age": 30})
@@ -165,8 +168,8 @@ def test_value_stored_as_flattened_object(self, store: ElasticsearchStore, es_cl
165168
{
166169
"collection": "test",
167170
"key": "test_key",
168-
"value": {"string": '{"name": "Alice", "age": 30}'},
169-
"created_at": "2025-10-27T15:33:31.795253+00:00",
171+
"value": {"flattened": {"name": "Alice", "age": 30}},
172+
"created_at": IsStr(min_length=20, max_length=40),
170173
}
171174
)
172175

@@ -177,9 +180,9 @@ def test_value_stored_as_flattened_object(self, store: ElasticsearchStore, es_cl
177180
{
178181
"collection": "test",
179182
"key": "test_key",
180-
"value": {"string": '{"name": "Bob", "age": 25}'},
181-
"created_at": "2025-10-27T15:33:32.040020+00:00",
182-
"expires_at": "2025-10-27T15:33:42.040020+00:00",
183+
"value": {"flattened": {"name": "Bob", "age": 25}},
184+
"created_at": IsStr(min_length=20, max_length=40),
185+
"expires_at": IsStr(min_length=20, max_length=40),
183186
}
184187
)
185188

@@ -198,13 +201,14 @@ def test_migration_from_legacy_mode(self, store: ElasticsearchStore, es_client:
198201
assert result == snapshot(None)
199202

200203

201-
class TestElasticsearchStoreNonNativeMode(TestElasticsearchStore):
204+
class TestElasticsearchStoreNonNativeMode(BaseTestElasticsearchStore):
205+
"""Test Elasticsearch store in non-native mode"""
206+
202207
@override
203208
@pytest.fixture
204-
def store(self) -> Generator[ElasticsearchStore, None, None]:
209+
def store(self) -> ElasticsearchStore:
205210
self._cleanup()
206-
with ElasticsearchStore(url=ES_URL, index_prefix="kv-store-e2e-test", native_storage=False) as store:
207-
yield store
211+
return ElasticsearchStore(url=ES_URL, index_prefix="kv-store-e2e-test", native_storage=False)
208212

209213
def test_value_stored_as_json_string(self, store: ElasticsearchStore, es_client: Elasticsearch):
210214
"""Verify values are stored as flattened objects, not JSON strings"""
@@ -220,8 +224,8 @@ def test_value_stored_as_json_string(self, store: ElasticsearchStore, es_client:
220224
{
221225
"collection": "test",
222226
"key": "test_key",
223-
"value": {"string": '{"name": "Alice", "age": 30}'},
224-
"created_at": "2025-10-27T15:33:32.631007+00:00",
227+
"value": {"string": '{"age": 30, "name": "Alice"}'},
228+
"created_at": IsStr(min_length=20, max_length=40),
225229
}
226230
)
227231

@@ -232,9 +236,9 @@ def test_value_stored_as_json_string(self, store: ElasticsearchStore, es_client:
232236
{
233237
"collection": "test",
234238
"key": "test_key",
235-
"value": {"string": '{"name": "Bob", "age": 25}'},
236-
"created_at": "2025-10-27T15:33:32.650061+00:00",
237-
"expires_at": "2025-10-27T15:33:42.650061+00:00",
239+
"value": {"string": '{"age": 25, "name": "Bob"}'},
240+
"created_at": IsStr(min_length=20, max_length=40),
241+
"expires_at": IsStr(min_length=20, max_length=40),
238242
}
239243
)
240244

0 commit comments

Comments
 (0)