All async helper tests passing

Author: sethmlarson

elasticsearch/_async/helpers/__init__.py

@@ -1,4 +1,3 @@
 # Licensed to Elasticsearch B.V under one or more agreements.
 # Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
 # See the LICENSE file in the project root for more information
-

test_elasticsearch/test_async/test_server/conftest.py

@@ -17,10 +17,9 @@ async def async_client():
 
     kw = {"timeout": 30, "ca_certs": ".ci/certs/ca.pem"}
     if "PYTHON_CONNECTION_CLASS" in os.environ:
-        from elasticsearch import connection
 
         kw["connection_class"] = getattr(
-            connection, os.environ["PYTHON_CONNECTION_CLASS"]
+            elasticsearch, os.environ["PYTHON_CONNECTION_CLASS"]
         )
 
     client = elasticsearch.AsyncElasticsearch(

test_elasticsearch/test_async/test_server/test_helpers.py

@@ -3,7 +3,7 @@
 # See the LICENSE file in the project root for more information
 
 import pytest
-from mock import patch
+from mock import patch, Mock
 
 from elasticsearch import helpers, TransportError
 from elasticsearch.helpers import ScanError
@@ -312,20 +312,29 @@ async def scan_fixture(async_client):
     await async_client.clear_scroll(scroll_id="_all")
 
 
-class TestScan:
-    mock_scroll_responses = [
-        {
-            "_scroll_id": "dummy_id",
-            "_shards": {"successful": 4, "total": 5, "skipped": 0},
-            "hits": {"hits": [{"scroll_data": 42}]},
-        },
-        {
-            "_scroll_id": "dummy_id",
-            "_shards": {"successful": 4, "total": 5, "skipped": 0},
-            "hits": {"hits": []},
-        },
-    ]
+class MockScroll:
+    def __init__(self):
+        self.i = 0
+        self.values = [
+            {
+                "_scroll_id": "dummy_id",
+                "_shards": {"successful": 4, "total": 5, "skipped": 0},
+                "hits": {"hits": [{"scroll_data": 42}]},
+            },
+            {
+                "_scroll_id": "dummy_id",
+                "_shards": {"successful": 4, "total": 5, "skipped": 0},
+                "hits": {"hits": []},
+            },
+        ]
+
+    async def scroll(self, *args, **kwargs):
+        val = self.values[self.i]
+        self.i += 1
+        return val
 
+
+class TestScan:
     async def test_order_can_be_preserved(self, async_client, scan_fixture):
         bulk = []
         for x in range(100):
@@ -373,7 +382,7 @@ async def test_scroll_error(self, async_client, scan_fixture):
         await async_client.bulk(bulk, refresh=True)
 
         with patch.object(async_client, "scroll") as scroll_mock:
-            scroll_mock.side_effect = self.mock_scroll_responses
+            scroll_mock.side_effect = MockScroll().scroll
             data = [
                 doc
                 async for doc in (
@@ -389,7 +398,7 @@ async def test_scroll_error(self, async_client, scan_fixture):
             assert len(data) == 3
             assert data[-1] == {"scroll_data": 42}
 
-            scroll_mock.side_effect = self.mock_scroll_responses
+            scroll_mock.side_effect = MockScroll().scroll
             with pytest.raises(ScanError):
                 data = [
                     doc
@@ -406,54 +415,62 @@ async def test_scroll_error(self, async_client, scan_fixture):
             assert len(data) == 3
             assert data[-1] == {"scroll_data": 42}
 
-    async def test_initial_search_error(self, async_client, scan_fixture):
-        with patch.object(self, "client") as client_mock:
-            client_mock.search.return_value = {
+    async def test_initial_search_error(self):
+        client_mock = Mock()
+
+        async def search_mock(*_, **__):
+            return {
                 "_scroll_id": "dummy_id",
                 "_shards": {"successful": 4, "total": 5, "skipped": 0},
                 "hits": {"hits": [{"search_data": 1}]},
             }
-            client_mock.scroll.side_effect = self.mock_scroll_responses
 
+        async def clear_scroll(*_, **__):
+            return {}
+
+        client_mock.search = search_mock
+        client_mock.scroll = MockScroll().scroll
+        client_mock.clear_scroll = clear_scroll
+
+        data = [
+            doc
+            async for doc in (
+                helpers.async_scan(
+                    client_mock, index="test_index", size=2, raise_on_error=False
+                )
+            )
+        ]
+        assert data == [{"search_data": 1}, {"scroll_data": 42}]
+
+        client_mock.scroll = MockScroll().scroll
+        with pytest.raises(ScanError):
             data = [
                 doc
                 async for doc in (
                     helpers.async_scan(
-                        async_client, index="test_index", size=2, raise_on_error=False
+                        client_mock, index="test_index", size=2, raise_on_error=True,
                     )
                 )
             ]
-            assert data == [{"search_data": 1}, {"scroll_data": 42}]
+            assert data == [{"search_data": 1}]
+            scroll_mock.assert_not_called()
 
-            client_mock.scroll.side_effect = self.mock_scroll_responses
-            with pytest.raises(ScanError):
-                data = [
-                    doc
-                    async for doc in (
-                        helpers.async_scan(
-                            async_client,
-                            index="test_index",
-                            size=2,
-                            raise_on_error=True,
-                        )
-                    )
-                ]
-                assert data == [{"search_data": 1}]
-                client_mock.scroll.assert_not_called()
+    async def test_no_scroll_id_fast_route(self):
+        client_mock = Mock()
 
-    async def test_no_scroll_id_fast_route(self, async_client, scan_fixture):
-        with patch.object(self, "client") as client_mock:
-            client_mock.search.return_value = {"no": "_scroll_id"}
-            data = [
-                doc
-                async for doc in (helpers.async_scan(async_client, index="test_index"))
-            ]
+        async def search_mock(*args, **kwargs):
+            return {"no": "_scroll_id"}
+
+        client_mock.search = search_mock
+        data = [
+            doc async for doc in (helpers.async_scan(client_mock, index="test_index"))
+        ]
 
-            assert data == []
-            client_mock.scroll.assert_not_called()
-            client_mock.clear_scroll.assert_not_called()
+        assert data == []
+        client_mock.scroll.assert_not_called()
+        client_mock.clear_scroll.assert_not_called()
 
-    @patch("elasticsearch.helpers.actions.logger")
+    @patch("elasticsearch._async.helpers.actions.logger")
     async def test_logger(self, logger_mock, async_client, scan_fixture):
         bulk = []
         for x in range(4):
@@ -462,7 +479,7 @@ async def test_logger(self, logger_mock, async_client, scan_fixture):
         await async_client.bulk(bulk, refresh=True)
 
         with patch.object(async_client, "scroll") as scroll_mock:
-            scroll_mock.side_effect = self.mock_scroll_responses
+            scroll_mock.side_effect = MockScroll().scroll
             _ = [
                 doc
                 async for doc in (
@@ -477,7 +494,7 @@ async def test_logger(self, logger_mock, async_client, scan_fixture):
             ]
             logger_mock.warning.assert_called()
 
-            scroll_mock.side_effect = self.mock_scroll_responses
+            scroll_mock.side_effect = MockScroll().scroll
             try:
                 _ = [
                     doc

test_elasticsearch/test_server/test_helpers.py

@@ -177,8 +177,6 @@ def test_rejected_documents_are_retried_at_most_max_retries_times(self):
         self.assertEqual({"value": 2, "relation": "eq"}, res["hits"]["total"])
         self.assertEqual(4, failing_client._called)
 
-        assert False
-
     def test_transport_error_is_raised_with_max_retries(self):
         failing_client = FailingBulkClient(
             self.client,