Add refresh parameter tests

elastic · Jun 29, 2023 · c23c975 · c23c975
1 parent 41b4c30
commit c23c975
Showing 1 changed file with 76 additions and 0 deletions.
diff --git a/tests/driver/runner_test.py b/tests/driver/runner_test.py
@@ -1243,6 +1243,82 @@ async def test_bulk_index_error_logs_warning_with_detailed_stats_body(self, es):
 
         es.bulk.assert_awaited_with(body=bulk_params["body"], params={})
 
+    @mock.patch("elasticsearch.Elasticsearch")
+    @pytest.mark.asyncio
+    async def test_bulk_index_success_with_refresh_true(self, es):
+        bulk_response = {
+            "errors": False,
+            "took": 8,
+            "items": [{"create": {"_index": "test", "result": "created", "status": 201, "forced_refresh": True}}],
+        }
+        es.bulk = mock.AsyncMock(return_value=bulk_response)
+
+        bulk = runner.BulkIndex()
+
+        bulk_params = {
+            "body": _build_bulk_body("index_line"),
+            "index": "test",
+            "action-metadata-present": False,
+            "type": "_doc",
+            "bulk-size": 1,
+            "unit": "docs",
+            "detailed-results": True,
+            "refresh": "true",
+        }
+
+        result = await bulk(es, dict(bulk_params))
+
+        assert result == {
+            "took": 8,
+            "index": "test",
+            "weight": 1,
+            "unit": "docs",
+            "success": True,
+            "success-count": 1,
+            "error-count": 0,
+            "bulk-request-size-bytes": 10,
+            "ops": {"create": collections.Counter({"item-count": 1, "created": 1})},
+            "shards_histogram": [],
+            "total-document-size-bytes": 10,
+        }
+
+        es.bulk.assert_awaited_with(doc_type="_doc", index="test", body=bulk_params["body"], params={"refresh": "true"})
+
+    @mock.patch("elasticsearch.Elasticsearch")
+    @pytest.mark.asyncio
+    async def test_bulk_index_success_with_refresh_wait_for(self, es):
+        bulk_response = {
+            "errors": False,
+            "took": 8,
+        }
+        es.bulk = mock.AsyncMock(return_value=io.BytesIO(json.dumps(bulk_response).encode()))
+
+        bulk = runner.BulkIndex()
+
+        bulk_params = {
+            "body": _build_bulk_body("index_line"),
+            "index": "test",
+            "action-metadata-present": False,
+            "type": "_doc",
+            "bulk-size": 1,
+            "unit": "docs",
+            "refresh": "wait_for",
+        }
+
+        result = await bulk(es, bulk_params)
+
+        assert result == {
+            "took": 8,
+            "index": "test",
+            "weight": 1,
+            "unit": "docs",
+            "success": True,
+            "success-count": 1,
+            "error-count": 0,
+        }
+
+        es.bulk.assert_awaited_with(doc_type="_doc", index="test", body=bulk_params["body"], params={"refresh": "wait_for"})
+
 
 class TestForceMergeRunner:
     @mock.patch("elasticsearch.Elasticsearch")