Prefer AsyncMock to as_future

tests/__init__.py

@@ -29,22 +29,3 @@ def async_wrapper(*args, **kwargs):
         asyncio.run(t(*args, **kwargs), debug=True)
 
     return async_wrapper
-
-
-def as_future(result=None, exception=None):
-    """
-
-    Helper to create a future that completes immediately either with a result or exceptionally.
-
-    :param result: Regular result.
-    :param exception: Exceptional result.
-    :return: The corresponding future.
-    """
-    f = asyncio.get_running_loop().create_future()
-    if exception and result:
-        raise AssertionError("Specify a result or an exception but not both")
-    if exception:
-        f.set_exception(exception)
-    else:
-        f.set_result(result)
-    return f

tests/driver/driver_test.py

@@ -29,7 +29,7 @@
 from esrally import config, exceptions, metrics, track
 from esrally.driver import driver, runner, scheduler
 from esrally.track import params
-from tests import as_future, run_async
+from tests import run_async
 
 
 class DriverTestParamSource:
@@ -1369,7 +1369,7 @@ async def test_execute_schedule_in_throughput_mode(self, es):
         task_start = time.perf_counter()
         es.new_request_context.return_value = AsyncExecutorTests.StaticRequestTiming(task_start=task_start)
 
-        es.bulk.return_value = as_future(io.StringIO('{"errors": false, "took": 8}'))
+        es.bulk = mock.AsyncMock(return_value=io.StringIO('{"errors": false, "took": 8}'))
 
         params.register_param_source_for_name("driver-test-param-source", DriverTestParamSource)
         test_track = track.Track(name="unittest", description="unittest track", indices=None, challenges=None)
@@ -1566,8 +1566,8 @@ async def test_execute_schedule_runner_overrides_times(self, es):
     @mock.patch("elasticsearch.Elasticsearch")
     @run_async
     async def test_execute_schedule_throughput_throttled(self, es):
-        def perform_request(*args, **kwargs):
-            return as_future()
+        async def perform_request(*args, **kwargs):
+            return None
 
         es.init_request_context.return_value = {"request_start": 0, "request_end": 10}
         # as this method is called several times we need to return a fresh instance every time as the previous
@@ -1635,7 +1635,7 @@ def perform_request(*args, **kwargs):
     @run_async
     async def test_cancel_execute_schedule(self, es):
         es.init_request_context.return_value = {"request_start": 0, "request_end": 10}
-        es.bulk.return_value = as_future(io.StringIO('{"errors": false, "took": 8}'))
+        es.bulk = mock.AsyncMock(return_value=io.StringIO('{"errors": false, "took": 8}'))
 
         params.register_param_source_for_name("driver-test-param-source", DriverTestParamSource)
         test_track = track.Track(name="unittest", description="unittest track", indices=None, challenges=None)
@@ -1742,8 +1742,7 @@ async def __call__(self):
     async def test_execute_single_no_return_value(self):
         es = None
         params = None
-        runner = mock.Mock()
-        runner.return_value = as_future()
+        runner = mock.AsyncMock()
 
         ops, unit, request_meta_data = await driver.execute_single(self.context_managed(runner), es, params, on_error="continue")
 
@@ -1755,8 +1754,7 @@ async def test_execute_single_no_return_value(self):
     async def test_execute_single_tuple(self):
         es = None
         params = None
-        runner = mock.Mock()
-        runner.return_value = as_future(result=(500, "MB"))
+        runner = mock.AsyncMock(return_value=(500, "MB"))
 
         ops, unit, request_meta_data = await driver.execute_single(self.context_managed(runner), es, params, on_error="continue")
 
@@ -1768,9 +1766,8 @@ async def test_execute_single_tuple(self):
     async def test_execute_single_dict(self):
         es = None
         params = None
-        runner = mock.Mock()
-        runner.return_value = as_future(
-            {
+        runner = mock.AsyncMock(
+            return_value={
                 "weight": 50,
                 "unit": "docs",
                 "some-custom-meta-data": "valid",
@@ -1798,7 +1795,7 @@ async def test_execute_single_with_connection_error_always_aborts(self):
                 es = None
                 params = None
                 # ES client uses pseudo-status "N/A" in this case...
-                runner = mock.Mock(side_effect=as_future(exception=elasticsearch.ConnectionError("N/A", "no route to host", None)))
+                runner = mock.AsyncMock(side_effect=elasticsearch.ConnectionError("N/A", "no route to host", None))
 
                 with self.assertRaises(exceptions.RallyAssertionError) as ctx:
                     await driver.execute_single(self.context_managed(runner), es, params, on_error=on_error)
@@ -1808,9 +1805,7 @@ async def test_execute_single_with_connection_error_always_aborts(self):
     async def test_execute_single_with_http_400_aborts_when_specified(self):
         es = None
         params = None
-        runner = mock.Mock(
-            side_effect=as_future(exception=elasticsearch.NotFoundError(404, "not found", "the requested document could not be found"))
-        )
+        runner = mock.AsyncMock(side_effect=elasticsearch.NotFoundError(404, "not found", "the requested document could not be found"))
 
         with self.assertRaises(exceptions.RallyAssertionError) as ctx:
             await driver.execute_single(self.context_managed(runner), es, params, on_error="abort")
@@ -1823,9 +1818,7 @@ async def test_execute_single_with_http_400_aborts_when_specified(self):
     async def test_execute_single_with_http_400(self):
         es = None
         params = None
-        runner = mock.Mock(
-            side_effect=as_future(exception=elasticsearch.NotFoundError(404, "not found", "the requested document could not be found"))
-        )
+        runner = mock.AsyncMock(side_effect=elasticsearch.NotFoundError(404, "not found", "the requested document could not be found"))
 
         ops, unit, request_meta_data = await driver.execute_single(self.context_managed(runner), es, params, on_error="continue")
 
@@ -1845,7 +1838,7 @@ async def test_execute_single_with_http_400(self):
     async def test_execute_single_with_http_413(self):
         es = None
         params = None
-        runner = mock.Mock(side_effect=as_future(exception=elasticsearch.NotFoundError(413, b"", b"")))
+        runner = mock.AsyncMock(side_effect=elasticsearch.NotFoundError(413, b"", b""))
 
         ops, unit, request_meta_data = await driver.execute_single(self.context_managed(runner), es, params, on_error="continue")