feat(api_core): add retry param into PollingFuture() and it's inheritors (#9923)
* feat(api_core): add retry param into PollingFuture() and it's inheritors
Towards #6197
diff --git a/google/api_core/future/polling.py b/google/api_core/future/polling.py
index 4266e9e..6b4c687 100644
--- a/google/api_core/future/polling.py
+++ b/google/api_core/future/polling.py
@@ -66,9 +66,12 @@
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 87f42a9..e6407b8 100644
--- a/google/api_core/operation.py
+++ b/google/api_core/operation.py
@@ -145,21 +145,28 @@
)
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 a5346a7..14b95cb 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 @@
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 = [