Support trio with httpx

Author: pquentin

elasticsearch/_async/helpers.py

@@ -15,7 +15,7 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-import asyncio
+import anyio
 import logging
 from typing import (
     Any,
@@ -245,7 +245,7 @@ async def map_actions() -> AsyncIterable[_TYPE_BULK_ACTION_HEADER_AND_BODY]:
                 ]
             ] = []
             if attempt:
-                await asyncio.sleep(
+                await anyio.sleep(
                     min(max_backoff, initial_backoff * 2 ** (attempt - 1))
                 )
 

pyproject.toml

@@ -41,11 +41,18 @@ keywords = [
 ]
 dynamic = ["version"]
 dependencies = [
-    "elastic-transport>=9.1.0,<10",
+    # TODO revert before merging/releasing
+    "elastic-transport @ git+https://github.com/pquentin/elastic-transport-python.git@trio-support",
     "python-dateutil",
     "typing-extensions",
+    "anyio",
+    "sniffio",
 ]
 
+# TODO revert before merging/releasing
+[tool.hatch.metadata]
+allow-direct-references = true
+
 [project.optional-dependencies]
 async = ["aiohttp>=3,<4"]
 requests = ["requests>=2.4.0, !=2.32.2, <3.0.0"]
@@ -56,6 +63,7 @@ vectorstore_mmr = ["numpy>=1", "simsimd>=3"]
 dev = [
     "requests>=2, <3",
     "aiohttp",
+    "httpx",
     "pytest",
     "pytest-cov",
     "pytest-mock",
@@ -78,6 +86,7 @@ dev = [
     "mapbox-vector-tile",
     "jinja2",
     "tqdm",
+    "trio",
     "mypy",
     "pyright",
     "types-python-dateutil",

test_elasticsearch/test_async/test_server/conftest.py

@@ -15,28 +15,29 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
+import sniffio
 import pytest
-import pytest_asyncio
 
 import elasticsearch
 
 from ...utils import CA_CERTS, wipe_cluster
 
-pytestmark = pytest.mark.asyncio
 
 
-@pytest_asyncio.fixture(scope="function")
+@pytest.fixture(scope="function")
 async def async_client_factory(elasticsearch_url):
-
-    if not hasattr(elasticsearch, "AsyncElasticsearch"):
-        pytest.skip("test requires 'AsyncElasticsearch' and aiohttp to be installed")
-
+    print("async!", elasticsearch_url)
+    kwargs = {}
+    if sniffio.current_async_library() == "trio":
+        kwargs["node_class"] = "httpxasync"
     # Unfortunately the asyncio client needs to be rebuilt every
     # test execution due to how pytest-asyncio manages
     # event loops (one per test!)
     client = None
     try:
-        client = elasticsearch.AsyncElasticsearch(elasticsearch_url, ca_certs=CA_CERTS)
+        client = elasticsearch.AsyncElasticsearch(
+            elasticsearch_url, ca_certs=CA_CERTS, **kwargs
+        )
         yield client
     finally:
         if client:

test_elasticsearch/test_async/test_server/test_clients.py

@@ -18,7 +18,7 @@
 
 import pytest
 
-pytestmark = pytest.mark.asyncio
+pytestmark = pytest.mark.anyio
 
 
 @pytest.mark.parametrize("kwargs", [{"body": {"text": "привет"}}, {"text": "привет"}])

test_elasticsearch/test_async/test_server/test_helpers.py

@@ -15,20 +15,19 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-import asyncio
+import anyio
 import logging
 from datetime import datetime, timedelta, timezone
 from unittest.mock import MagicMock, call, patch
 
 import pytest
-import pytest_asyncio
 from elastic_transport import ApiResponseMeta, ObjectApiResponse
 
 from elasticsearch import helpers
 from elasticsearch.exceptions import ApiError
 from elasticsearch.helpers import ScanError
 
-pytestmark = [pytest.mark.asyncio]
+pytestmark = pytest.mark.anyio
 
 
 class AsyncMock(MagicMock):
@@ -92,7 +91,7 @@ async def test_all_documents_get_inserted(self, async_client):
     async def test_documents_data_types(self, async_client):
         async def async_gen():
             for x in range(100):
-                await asyncio.sleep(0)
+                await anyio.sleep(0)
                 yield {"answer": x, "_id": x}
 
         def sync_gen():
@@ -491,7 +490,7 @@ def __await__(self):
         return self().__await__()
 
 
-@pytest_asyncio.fixture(scope="function")
+@pytest.fixture(scope="function")
 async def scan_teardown(async_client):
     yield
     await async_client.clear_scroll(scroll_id="_all")
@@ -915,7 +914,7 @@ async def test_scan_from_keyword_is_aliased(async_client, scan_kwargs):
         assert "from" not in search_mock.call_args[1]
 
 
-@pytest_asyncio.fixture(scope="function")
+@pytest.fixture(scope="function")
 async def reindex_setup(async_client):
     bulk = []
     for x in range(100):
@@ -993,7 +992,7 @@ async def test_all_documents_get_moved(self, async_client, reindex_setup):
         )["_source"]
 
 
-@pytest_asyncio.fixture(scope="function")
+@pytest.fixture(scope="function")
 async def parent_reindex_setup(async_client):
     body = {
         "settings": {"number_of_shards": 1, "number_of_replicas": 0},
@@ -1054,7 +1053,7 @@ async def test_children_are_reindexed_correctly(
         } == q
 
 
-@pytest_asyncio.fixture(scope="function")
+@pytest.fixture(scope="function")
 async def reindex_data_stream_setup(async_client):
     dt = datetime.now(tz=timezone.utc)
     bulk = []

test_elasticsearch/test_async/test_server/test_mapbox_vector_tile.py

@@ -16,14 +16,13 @@
 #  under the License.
 
 import pytest
-import pytest_asyncio
 
 from elasticsearch import RequestError
 
-pytestmark = pytest.mark.asyncio
+pytestmark = pytest.mark.anyio
 
 
-@pytest_asyncio.fixture(scope="function")
+@pytest.fixture(scope="function")
 async def mvt_setup(async_client):
     await async_client.indices.create(
         index="museums",

test_elasticsearch/test_async/test_server/test_rest_api_spec.py

@@ -25,7 +25,6 @@
 import warnings
 
 import pytest
-import pytest_asyncio
 
 from elasticsearch import ElasticsearchWarning, RequestError
 
@@ -39,7 +38,7 @@
 )
 from ...utils import parse_version
 
-pytestmark = pytest.mark.asyncio
+pytestmark = pytest.mark.anyio
 
 XPACK_FEATURES = None
 ES_VERSION = None
@@ -240,7 +239,7 @@ async def _feature_enabled(self, name):
         return name in XPACK_FEATURES
 
 
-@pytest_asyncio.fixture(scope="function")
+@pytest.fixture(scope="function")
 def async_runner(async_client_factory):
     return AsyncYamlRunner(async_client_factory)
 

test_elasticsearch/test_async/test_transport.py

@@ -16,7 +16,8 @@
 #  under the License.
 
 
-import asyncio
+import time
+import anyio
 import re
 import warnings
 from typing import Any, Dict, Optional
@@ -40,8 +41,6 @@
     UnsupportedProductError,
 )
 
-pytestmark = pytest.mark.asyncio
-
 
 class DummyNode(BaseAsyncNode):
     def __init__(self, config: NodeConfig):
@@ -175,6 +174,7 @@ def mark_live(self, connection):
 }"""
 
 
+@pytest.mark.anyio
 class TestTransport:
     async def test_request_timeout_extracted_from_params_and_passed(self):
         client = AsyncElasticsearch(
@@ -378,6 +378,9 @@ async def test_override_mark_dead_mark_live(self):
         assert len(client.transport.node_pool._alive_nodes) == 2
         assert len(client.transport.node_pool._dead_consecutive_failures) == 0
 
+
+@pytest.mark.asyncio
+class TestSniffing:
     @pytest.mark.parametrize(
         ["nodes_info_response", "node_host"],
         [(CLUSTER_NODES, "1.1.1.1"), (CLUSTER_NODES_7x_PUBLISH_HOST, "somehost.tld")],
@@ -528,23 +531,22 @@ async def test_sniff_on_node_failure_triggers(self, extra_key, extra_value):
         assert len(client.transport.node_pool) == 3
 
     async def test_sniff_after_n_seconds(self):
-        event_loop = asyncio.get_running_loop()
         client = AsyncElasticsearch(  # noqa: F821
             [NodeConfig("http", "localhost", 9200, _extras={"data": CLUSTER_NODES})],
             node_class=DummyNode,
             min_delay_between_sniffing=5,
         )
-        client.transport._last_sniffed_at = event_loop.time()
+        client.transport._last_sniffed_at = time.monotonic()
 
         await client.info()
 
         for _ in range(4):
             await client.info()
-            await asyncio.sleep(0)
+            await anyio.sleep(0)
 
         assert 1 == len(client.transport.node_pool)
 
-        client.transport._last_sniffed_at = event_loop.time() - 5.1
+        client.transport._last_sniffed_at = time.monotonic() - 5.1
 
         await client.info()
         await client.transport._sniffing_task  # Need to wait for the sniffing task to complete
@@ -554,9 +556,9 @@ async def test_sniff_after_n_seconds(self):
             node.base_url for node in client.transport.node_pool.all()
         )
         assert (
-            event_loop.time() - 1
+            time.monotonic() - 1
             < client.transport._last_sniffed_at
-            < event_loop.time() + 0.01
+            < time.monotonic() + 0.01
         )
 
     @pytest.mark.parametrize(
@@ -580,8 +582,10 @@ async def test_sniffing_disabled_on_elastic_cloud(self, kwargs):
             == "Sniffing should not be enabled when connecting to Elastic Cloud"
         )
 
-    async def test_sniff_on_start_close_unlocks_async_calls(self):
-        event_loop = asyncio.get_running_loop()
+    async def test_sniff_on_start_close_unlocks_async_calls(self, anyio_backend):
+        if anyio_backend == "trio":
+            pytest.skip("trio does not support sniffing")
+
         client = AsyncElasticsearch(  # noqa: F821
             [
                 NodeConfig(
@@ -596,20 +600,17 @@ async def test_sniff_on_start_close_unlocks_async_calls(self):
         )
 
         # Start making _async_calls() before we cancel
-        tasks = []
-        start_time = event_loop.time()
-        for _ in range(3):
-            tasks.append(event_loop.create_task(client.info()))
-            await asyncio.sleep(0)
-
-        # Close the transport while the sniffing task is active! :(
-        await client.transport.close()
-
-        # Now we start waiting on all those _async_calls()
-        await asyncio.gather(*tasks)
-        end_time = event_loop.time()
-        duration = end_time - start_time
+        async with anyio.create_task_group() as tg:
+            start_time = time.monotonic()
+            for _ in range(3):
+                tg.start_soon(client.info)
+            await anyio.sleep(0)
+
+            # Close the transport while the sniffing task is active! :(
+            await client.transport.close()
 
+        end_time = time.monotonic()
+        duration = end_time - start_time
         # A lot quicker than 10 seconds defined in 'delay'
         assert duration < 1
 
@@ -661,6 +662,7 @@ def sniffed_node_callback(
         assert ports == {9200, 124}
 
 
+@pytest.mark.anyio
 @pytest.mark.parametrize("headers", [{}, {"X-elastic-product": "BAD HEADER"}])
 async def test_unsupported_product_error(headers):
     client = AsyncElasticsearch(
@@ -690,6 +692,7 @@ async def test_unsupported_product_error(headers):
     )
 
 
+@pytest.mark.anyio
 @pytest.mark.parametrize("status", [401, 403, 413, 500])
 async def test_unsupported_product_error_not_raised_on_non_2xx(status):
     client = AsyncElasticsearch(
@@ -709,6 +712,7 @@ async def test_unsupported_product_error_not_raised_on_non_2xx(status):
         assert e.meta.status == status
 
 
+@pytest.mark.anyio
 @pytest.mark.parametrize("status", [404, 500])
 async def test_api_error_raised_before_product_error(status):
     client = AsyncElasticsearch(
@@ -737,6 +741,7 @@ async def test_api_error_raised_before_product_error(status):
     assert calls[0][0] == ("GET", "/")
 
 
+@pytest.mark.anyio
 @pytest.mark.parametrize(
     "headers",
     [

test_elasticsearch/test_dsl/_async/test_document.py

@@ -582,21 +582,21 @@ def test_meta_fields_can_be_set_directly_in_init() -> None:
     assert md.meta.id is p
 
 
-@pytest.mark.asyncio
+@pytest.mark.anyio
 async def test_save_no_index(async_mock_client: Any) -> None:
     md = MyDoc()
     with raises(ValidationException):
         await md.save(using="mock")
 
 
-@pytest.mark.asyncio
+@pytest.mark.anyio
 async def test_delete_no_index(async_mock_client: Any) -> None:
     md = MyDoc()
     with raises(ValidationException):
         await md.delete(using="mock")
 
 
-@pytest.mark.asyncio
+@pytest.mark.anyio
 async def test_update_no_fields() -> None:
     md = MyDoc()
     with raises(IllegalOperation):

test_elasticsearch/test_dsl/_async/test_index.py

@@ -190,7 +190,7 @@ def test_index_template_can_have_order() -> None:
     assert {"index_patterns": ["i-*"], "order": 2} == it.to_dict()
 
 
-@pytest.mark.asyncio
+@pytest.mark.anyio
 async def test_index_template_save_result(async_mock_client: Any) -> None:
     it = AsyncIndexTemplate("test-template", "test-*")