integration tests

Author: miguelgrinberg

elasticsearch/_async/helpers.py

@@ -31,6 +31,7 @@
     Tuple,
     TypeVar,
     Union,
+    cast,
 )
 
 from ..compat import safe_task
@@ -41,6 +42,7 @@
     _TYPE_BULK_ACTION_HEADER,
     _TYPE_BULK_ACTION_HEADER_AND_BODY,
     _TYPE_BULK_ACTION_HEADER_WITH_META_AND_BODY,
+    _TYPE_BULK_ACTION_WITH_META,
     BulkMeta,
     _ActionChunker,
     _process_bulk_chunk_error,
@@ -195,7 +197,10 @@ async def azip(
 
 async def async_streaming_bulk(
     client: AsyncElasticsearch,
-    actions: Union[Iterable[_TYPE_BULK_ACTION], AsyncIterable[_TYPE_BULK_ACTION]],
+    actions: Union[
+        Iterable[_TYPE_BULK_ACTION_WITH_META],
+        AsyncIterable[_TYPE_BULK_ACTION_WITH_META],
+    ],
     chunk_size: int = 500,
     max_chunk_bytes: int = 100 * 1024 * 1024,
     flush_after_seconds: Optional[float] = None,
@@ -257,9 +262,14 @@ async def async_streaming_bulk(
     if isinstance(retry_on_status, int):
         retry_on_status = (retry_on_status,)
 
-    async def map_actions() -> AsyncIterable[_TYPE_BULK_ACTION_HEADER_AND_BODY]:
+    async def map_actions() -> (
+        AsyncIterable[_TYPE_BULK_ACTION_HEADER_WITH_META_AND_BODY]
+    ):
         async for item in aiter(actions):
-            yield expand_action_callback(item)
+            if isinstance(item, BulkMeta):
+                yield item, None
+            else:
+                yield expand_action_callback(cast(_TYPE_BULK_ACTION, item))
 
     serializer = client.transport.serializers.get_serializer("application/json")
 

elasticsearch/compat.py

@@ -22,7 +22,7 @@
 from contextlib import asynccontextmanager, contextmanager
 from pathlib import Path
 from threading import Thread
-from typing import Tuple, Type, Union
+from typing import Any, AsyncIterator, Callable, Coroutine, Iterator, Tuple, Type, Union
 
 string_types: Tuple[Type[str], Type[bytes]] = (str, bytes)
 
@@ -80,15 +80,17 @@ def warn_stacklevel() -> int:
 
 
 @contextmanager
-def safe_thread(target, *args, **kwargs):
+def safe_thread(
+    target: Callable[..., Any], *args: Any, **kwargs: Any
+) -> Iterator[Thread]:
     """Run a thread within a context manager block.
 
     The thread is automatically joined when the block ends. If the thread raised
     an exception, it is raised in the caller's context.
     """
     captured_exception = None
 
-    def run():
+    def run() -> None:
         try:
             target(*args, **kwargs)
         except BaseException as exc:
@@ -97,14 +99,14 @@ def run():
 
     thread = Thread(target=run)
     thread.start()
-    yield
+    yield thread
     thread.join()
     if captured_exception:
         raise captured_exception
 
 
 @asynccontextmanager
-async def safe_task(coro):
+async def safe_task(coro: Coroutine[Any, Any, Any]) -> AsyncIterator[asyncio.Task[Any]]:
     """Run a background task within a context manager block.
 
     The task is awaited when the block ends.

elasticsearch/helpers/actions.py

@@ -20,7 +20,6 @@
 import time
 from enum import Enum
 from operator import methodcaller
-from threading import Thread
 from typing import (
     Any,
     Callable,

test_elasticsearch/test_async/test_server/test_helpers.py

@@ -17,6 +17,7 @@
 
 import asyncio
 import logging
+import time
 from datetime import datetime, timedelta, timezone
 from unittest.mock import MagicMock, call, patch
 
@@ -123,6 +124,49 @@ def sync_gen():
             "_source"
         ]
 
+    async def test_explicit_flushes(self, async_client):
+        async def async_gen():
+            for x in range(2):
+                yield {"answer": x, "_id": x}
+            yield helpers.BULK_FLUSH
+            await asyncio.sleep(0.1)
+            for x in range(2):
+                yield {"answer": x + 2, "_id": x + 2}
+
+        timestamps = []
+        async for ok, item in helpers.async_streaming_bulk(
+            async_client, async_gen(), index="test-index", refresh=True
+        ):
+            assert ok
+            timestamps.append(time.time())
+
+        assert timestamps[1] - timestamps[0] < 0.1
+        assert timestamps[2] - timestamps[1] > 0.1
+        assert timestamps[3] - timestamps[2] < 0.1
+
+    async def test_timeout_flushes(self, async_client):
+        async def async_gen():
+            for x in range(2):
+                yield {"answer": x, "_id": x}
+            await asyncio.sleep(0.5)
+            for x in range(2):
+                yield {"answer": x + 2, "_id": x + 2}
+
+        timestamps = []
+        async for ok, item in helpers.async_streaming_bulk(
+            async_client,
+            async_gen(),
+            index="test-index",
+            refresh=True,
+            flush_after_seconds=0.05,
+        ):
+            assert ok
+            timestamps.append(time.time())
+
+        assert timestamps[1] - timestamps[0] < 0.1
+        assert timestamps[2] - timestamps[1] > 0.1
+        assert timestamps[3] - timestamps[2] < 0.1
+
     async def test_all_errors_from_chunk_are_raised_on_failure(self, async_client):
         await async_client.indices.create(
             index="i",

test_elasticsearch/test_server/test_helpers.py

@@ -16,6 +16,7 @@
 #  under the License.
 
 import json
+import time
 from datetime import datetime, timedelta
 from unittest.mock import call, patch
 
@@ -75,6 +76,51 @@ def test_bulk_all_documents_get_inserted(sync_client):
     assert {"answer": 42} == sync_client.get(index="test-index", id=42)["_source"]
 
 
+def test_explicit_flushes(sync_client):
+    def sync_gen():
+        for x in range(2):
+            yield {"answer": x, "_id": x}
+        yield helpers.BULK_FLUSH
+        time.sleep(0.1)
+        for x in range(2):
+            yield {"answer": x + 2, "_id": x + 2}
+
+    timestamps = []
+    for ok, item in helpers.streaming_bulk(
+        sync_client, sync_gen(), index="test-index", refresh=True
+    ):
+        assert ok
+        timestamps.append(time.time())
+
+    assert timestamps[1] - timestamps[0] < 0.1
+    assert timestamps[2] - timestamps[1] > 0.1
+    assert timestamps[3] - timestamps[2] < 0.1
+
+
+def test_timeout_flushes(sync_client):
+    def sync_gen():
+        for x in range(2):
+            yield {"answer": x, "_id": x}
+        time.sleep(0.5)
+        for x in range(2):
+            yield {"answer": x + 2, "_id": x + 2}
+
+    timestamps = []
+    for ok, item in helpers.streaming_bulk(
+        sync_client,
+        sync_gen(),
+        index="test-index",
+        refresh=True,
+        flush_after_seconds=0.05,
+    ):
+        assert ok
+        timestamps.append(time.time())
+
+    assert timestamps[1] - timestamps[0] < 0.1
+    assert timestamps[2] - timestamps[1] > 0.1
+    assert timestamps[3] - timestamps[2] < 0.1
+
+
 def test_bulk_all_errors_from_chunk_are_raised_on_failure(sync_client):
     sync_client.indices.create(
         index="i",