unit tests

Author: miguelgrinberg

elasticsearch/helpers/__init__.py

@@ -19,12 +19,21 @@
 from .._utils import fixup_module_metadata
 from .actions import _chunk_actions  # noqa: F401
 from .actions import _process_bulk_chunk  # noqa: F401
-from .actions import bulk, expand_action, parallel_bulk, reindex, scan, streaming_bulk
+from .actions import (
+    BULK_FLUSH,
+    bulk,
+    expand_action,
+    parallel_bulk,
+    reindex,
+    scan,
+    streaming_bulk,
+)
 from .errors import BulkIndexError, ScanError
 
 __all__ = [
     "BulkIndexError",
     "ScanError",
+    "BULK_FLUSH",
     "expand_action",
     "streaming_bulk",
     "bulk",

elasticsearch/helpers/actions.py

@@ -50,13 +50,14 @@ class BulkMeta(Enum):
     flush = 1
 
 
-FLUSH_NOW = BulkMeta.flush
+BULK_FLUSH = BulkMeta.flush
 
 _TYPE_BULK_ACTION = Union[bytes, str, Dict[str, Any], BulkMeta]
-_TYPE_BULK_ACTION_HEADER = Union[Dict[str, Any], BulkMeta]
+_TYPE_BULK_ACTION_HEADER = Dict[str, Any]
+_TYPE_BULK_ACTION_HEADER_AND_META = Union[Dict[str, Any], BulkMeta]
 _TYPE_BULK_ACTION_BODY = Union[None, bytes, Dict[str, Any]]
 _TYPE_BULK_ACTION_HEADER_AND_BODY = Tuple[
-    _TYPE_BULK_ACTION_HEADER, _TYPE_BULK_ACTION_BODY
+    _TYPE_BULK_ACTION_HEADER_AND_META, _TYPE_BULK_ACTION_BODY
 ]
 
 
@@ -150,7 +151,7 @@ def __init__(
         ] = []
 
     def feed(
-        self, action: _TYPE_BULK_ACTION_HEADER, data: _TYPE_BULK_ACTION_BODY
+        self, action: _TYPE_BULK_ACTION_HEADER_AND_META, data: _TYPE_BULK_ACTION_BODY
     ) -> Optional[
         Tuple[
             List[
@@ -163,17 +164,14 @@ def feed(
         ]
     ]:
         ret = None
-        raw_action = action
-        raw_data = data
-        action_bytes = None
-        data_bytes = None
+        action_bytes = b""
+        data_bytes: Optional[bytes] = None
         cur_size = 0
         if not isinstance(action, BulkMeta):
             action_bytes = to_bytes(self.serializer.dumps(action), "utf-8")
             # +1 to account for the trailing new line character
             cur_size = len(action_bytes) + 1
 
-            data_bytes: Optional[bytes]
             if data is not None:
                 data_bytes = to_bytes(self.serializer.dumps(data), "utf-8")
                 cur_size += len(data_bytes) + 1
@@ -192,13 +190,13 @@ def feed(
             self.size = 0
             self.action_count = 0
 
-        if action_bytes is not None:
+        if not isinstance(action, BulkMeta):
             self.bulk_actions.append(action_bytes)
             if data_bytes is not None:
                 self.bulk_actions.append(data_bytes)
-                self.bulk_data.append((raw_action, raw_data))
+                self.bulk_data.append((action, data))
             else:
-                self.bulk_data.append((raw_action,))
+                self.bulk_data.append((action,))
 
             self.size += cur_size
             self.action_count += 1

test_elasticsearch/test_helpers.py

@@ -85,6 +85,7 @@ def setup_method(self, _):
     def test_expand_action(self):
         assert helpers.expand_action({}) == ({"index": {}}, {})
         assert helpers.expand_action({"key": "val"}) == ({"index": {}}, {"key": "val"})
+        assert helpers.expand_action(helpers.BULK_FLUSH) == (helpers.BULK_FLUSH, {})
 
     def test_expand_action_actions(self):
         assert helpers.expand_action(
@@ -178,6 +179,25 @@ def test_chunks_are_chopped_by_byte_size_properly(self):
             chunk = b"".join(chunk_actions)
             assert len(chunk) <= max_byte_size
 
+    def test_chunks_are_chopped_by_flush(self):
+        flush = helpers.expand_action(helpers.BULK_FLUSH)
+        actions = (
+            self.actions[:3]
+            + [flush] * 2  # two consecutive flushes after 3 items
+            + self.actions[3:4]
+            + [flush]  # flush after one more item
+            + self.actions[4:]
+            + [flush]  # flush at the end
+        )
+        chunks = list(helpers._chunk_actions(actions, 100, 99999999, JSONSerializer()))
+        assert 3 == len(chunks)
+        assert len(chunks[0][0]) == 3
+        assert len(chunks[0][1]) == 6
+        assert len(chunks[1][0]) == 1
+        assert len(chunks[1][1]) == 2
+        assert len(chunks[2][0]) == 96
+        assert len(chunks[2][1]) == 192
+
 
 class TestExpandActions:
     @pytest.mark.parametrize("action", ["whatever", b"whatever"])