diff --git a/google/api_core/future/polling.py b/google/api_core/future/polling.py index 4266e9e7..6b4c687d 100644 --- a/google/api_core/future/polling.py +++ b/google/api_core/future/polling.py @@ -66,9 +66,12 @@ def __init__(self, retry=DEFAULT_RETRY): self._done_callbacks = [] @abc.abstractmethod - def done(self): + def done(self, retry=DEFAULT_RETRY): """Checks to see if the operation is complete. + Args: + retry (google.api_core.retry.Retry): (Optional) How to retry the RPC. + Returns: bool: True if the operation is complete, False otherwise. """ diff --git a/google/api_core/operation.py b/google/api_core/operation.py index 87f42a97..e6407b8c 100644 --- a/google/api_core/operation.py +++ b/google/api_core/operation.py @@ -145,21 +145,28 @@ def _set_result_from_operation(self): ) self.set_exception(exception) - def _refresh_and_update(self): - """Refresh the operation and update the result if needed.""" + def _refresh_and_update(self, retry=polling.DEFAULT_RETRY): + """Refresh the operation and update the result if needed. + + Args: + retry (google.api_core.retry.Retry): (Optional) How to retry the RPC. + """ # If the currently cached operation is done, no need to make another # RPC as it will not change once done. if not self._operation.done: - self._operation = self._refresh() + self._operation = self._refresh(retry=retry) self._set_result_from_operation() - def done(self): + def done(self, retry=polling.DEFAULT_RETRY): """Checks to see if the operation is complete. + Args: + retry (google.api_core.retry.Retry): (Optional) How to retry the RPC. + Returns: bool: True if the operation is complete, False otherwise. """ - self._refresh_and_update() + self._refresh_and_update(retry) return self._operation.done def cancel(self): diff --git a/tests/unit/test_operation.py b/tests/unit/test_operation.py index a5346a70..14b95cbb 100644 --- a/tests/unit/test_operation.py +++ b/tests/unit/test_operation.py @@ -15,8 +15,10 @@ import mock +from google.api_core import exceptions from google.api_core import operation from google.api_core import operations_v1 +from google.api_core import retry from google.longrunning import operations_pb2 from google.protobuf import struct_pb2 from google.rpc import code_pb2 @@ -113,6 +115,23 @@ def test_result(): assert future.done() +def test_done_w_retry(): + RETRY_PREDICATE = retry.if_exception_type(exceptions.TooManyRequests) + test_retry = retry.Retry(predicate=RETRY_PREDICATE) + + expected_result = struct_pb2.Struct() + responses = [ + make_operation_proto(), + # Second operation response includes the result. + make_operation_proto(done=True, response=expected_result), + ] + future, _, _ = make_operation_future(responses) + future._refresh = mock.Mock() + + future.done(retry=test_retry) + future._refresh.assert_called_once_with(retry=test_retry) + + def test_exception(): expected_exception = status_pb2.Status(message="meep") responses = [