fix(google-api-core): report background polling failures

packages/google-api-core/google/api_core/future/async_future.py

@@ -134,20 +134,31 @@ 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.
 
         If the operation is completed, the callback will be scheduled onto the
         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)
 

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()

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()