Fixes for Elasticsearch tests (#234)

Author: strawgate

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

@@ -1,5 +1,6 @@
 from collections.abc import AsyncGenerator
 from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any
 
 import pytest
 from dirty_equals import IsFloat, IsStr
@@ -19,6 +20,9 @@
 from tests.conftest import docker_container, should_skip_docker_tests
 from tests.stores.base import BaseStoreTests, ContextManagerStoreTestMixin
 
+if TYPE_CHECKING:
+    from elastic_transport._response import ObjectApiResponse
+
 TEST_SIZE_LIMIT = 1 * 1024 * 1024  # 1MB
 ES_HOST = "localhost"
 ES_PORT = 9200
@@ -41,7 +45,12 @@ async def ping_elasticsearch() -> bool:
     es_client: AsyncElasticsearch = get_elasticsearch_client()
 
     async with es_client:
-        return await es_client.ping()
+        if not await es_client.ping():
+            return False
+
+        status: ObjectApiResponse[dict[str, Any]] = await es_client.options(ignore_status=404).cluster.health(wait_for_status="green")
+
+        return status.body.get("status") == "green"
 
 
 async def cleanup_elasticsearch_indices(elasticsearch_client: AsyncElasticsearch):
@@ -102,10 +111,7 @@ async def setup_elasticsearch(self, request: pytest.FixtureRequest) -> AsyncGene
     @pytest.fixture
     async def es_client(self) -> AsyncGenerator[AsyncElasticsearch, None]:
         async with AsyncElasticsearch(hosts=[ES_URL]) as es_client:
-            try:
-                yield es_client
-            finally:
-                await es_client.close()
+            yield es_client
 
     @override
     @pytest.fixture

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

@@ -3,6 +3,7 @@
 # DO NOT CHANGE! Change the original file instead.
 from collections.abc import Generator
 from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING, Any
 
 import pytest
 from dirty_equals import IsFloat, IsStr
@@ -22,6 +23,9 @@
 from tests.code_gen.conftest import docker_container, should_skip_docker_tests
 from tests.code_gen.stores.base import BaseStoreTests, ContextManagerStoreTestMixin
 
+if TYPE_CHECKING:
+    from elastic_transport._response import ObjectApiResponse
+
 TEST_SIZE_LIMIT = 1 * 1024 * 1024  # 1MB
 ES_HOST = "localhost"
 ES_PORT = 9200
@@ -42,7 +46,12 @@ def ping_elasticsearch() -> bool:
     es_client: Elasticsearch = get_elasticsearch_client()
 
     with es_client:
-        return es_client.ping()
+        if not es_client.ping():
+            return False
+
+        status: ObjectApiResponse[dict[str, Any]] = es_client.options(ignore_status=404).cluster.health(wait_for_status="green")
+
+        return status.body.get("status") == "green"
 
 
 def cleanup_elasticsearch_indices(elasticsearch_client: Elasticsearch):
@@ -103,10 +112,7 @@ def setup_elasticsearch(self, request: pytest.FixtureRequest) -> Generator[None,
     @pytest.fixture
     def es_client(self) -> Generator[Elasticsearch, None, None]:
         with Elasticsearch(hosts=[ES_URL]) as es_client:
-            try:
-                yield es_client
-            finally:
-                es_client.close()
+            yield es_client
 
     @override
     @pytest.fixture