diff --git a/packages/google-api-core/google/api_core/future/async_future.py b/packages/google-api-core/google/api_core/future/async_future.py index 57fdd083a1c0..59a902106c67 100644 --- a/packages/google-api-core/google/api_core/future/async_future.py +++ b/packages/google-api-core/google/api_core/future/async_future.py @@ -134,6 +134,14 @@ async def exception(self, timeout=None): await self._blocking_poll(timeout=timeout) return self._future.exception() + async def _background_poll(self): + """Poll in the background and report failures through the future.""" + try: + await self._blocking_poll() + except Exception as exc: + if not self._future.done(): + self.set_exception(exc) + def add_done_callback(self, fn): """Add a callback to be executed when the operation is complete. @@ -141,13 +149,16 @@ def add_done_callback(self, fn): event loop. Otherwise, the callback will be stored and invoked when the future is done. + A background polling failure is stored as the future's exception and + also causes the callbacks to be invoked. + Args: fn (Callable[Future]): The callback to execute when the operation is complete. """ if self._background_task is None: self._background_task = asyncio.get_event_loop().create_task( - self._blocking_poll() + self._background_poll() ) self._future.add_done_callback(fn) diff --git a/packages/google-api-core/tests/asyncio/future/test_async_future.py b/packages/google-api-core/tests/asyncio/future/test_async_future.py index 659f41cf90ec..8e9f2ad90e5c 100644 --- a/packages/google-api-core/tests/asyncio/future/test_async_future.py +++ b/packages/google-api-core/tests/asyncio/future/test_async_future.py @@ -17,6 +17,7 @@ import pytest +# isort: split from google.api_core import exceptions from google.api_core.future import async_future @@ -225,3 +226,76 @@ def callback2(unused_future): assert future.poll_count == 1 await callback_called.wait() await callback_called2.wait() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error_type", [exceptions.Forbidden, ValueError]) +async def test_callback_polling_error(error_type): + future = AsyncFuture() + error = error_type("polling failed") + future.done = mock.AsyncMock(side_effect=error) + callback = mock.Mock() + + future.add_done_callback(callback) + await asyncio.gather(future._background_task, return_exceptions=True) + await asyncio.sleep(0) + + callback.assert_called_once_with(future._future) + assert await future.exception() is error + with pytest.raises(error_type) as exc_info: + await future.result() + assert exc_info.value is error + future.done.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_background_polling_error_after_result(): + future = AsyncFuture() + polling_started = asyncio.Event() + finish_polling = asyncio.Event() + + async def done(): + polling_started.set() + await finish_polling.wait() + raise ValueError("polling failed after completion") + + future.done = done + callback = mock.Mock() + future.add_done_callback(callback) + await polling_started.wait() + future.set_result(42) + finish_polling.set() + outcomes = await asyncio.gather(future._background_task, return_exceptions=True) + await asyncio.sleep(0) + + assert outcomes == [None] + callback.assert_called_once_with(future._future) + assert await future.result() == 42 + + +@pytest.mark.asyncio +async def test_background_polling_cancelled(): + future = AsyncFutureWithPoll() + callback = mock.Mock() + future.add_done_callback(callback) + await asyncio.sleep(0) + future._background_task.cancel() + + with pytest.raises(asyncio.CancelledError): + await future._background_task + + callback.assert_not_called() + assert not future._future.done() + + +@pytest.mark.asyncio +async def test_foreground_polling_error(): + future = AsyncFuture() + error = exceptions.Forbidden("polling failed") + future.done = mock.AsyncMock(side_effect=error) + + with pytest.raises(exceptions.Forbidden) as exc_info: + await future.result() + + assert exc_info.value is error + assert not future._future.done() diff --git a/packages/google-api-core/tests/asyncio/test_operation_async.py b/packages/google-api-core/tests/asyncio/test_operation_async.py index 601c8995d4f6..71a33e5befea 100644 --- a/packages/google-api-core/tests/asyncio/test_operation_async.py +++ b/packages/google-api-core/tests/asyncio/test_operation_async.py @@ -13,6 +13,8 @@ # limitations under the License. +import asyncio + import pytest try: @@ -30,6 +32,7 @@ from google.protobuf import struct_pb2 from google.rpc import code_pb2, status_pb2 +# isort: split from google.api_core import exceptions, operation_async, operations_v1, retry_async TEST_OPERATION_NAME = "test/operation" @@ -81,6 +84,25 @@ async def test_constructor(): assert await future.running() +@pytest.mark.asyncio +async def test_callback_refresh_error(): + future, refresh, _ = make_operation_future() + error = exceptions.Forbidden("polling permission denied") + refresh.side_effect = error + callback = mock.Mock() + + future.add_done_callback(callback) + await asyncio.gather(future._background_task, return_exceptions=True) + await asyncio.sleep(0) + + callback.assert_called_once_with(future._future) + assert await future.exception() is error + with pytest.raises(exceptions.Forbidden) as exc_info: + await future.result() + assert exc_info.value is error + refresh.assert_awaited_once() + + @pytest.mark.asyncio async def test_metadata(): expected_metadata = struct_pb2.Struct()