feat(api_core): add retry param into PollingFuture() and it's inheritors (#9923)

Gurov Ilya · busunkim96 · commit 8e7b7f8a5f4b · 2020-01-24T10:26:24.000-08:00
* feat(api_core): add retry param into PollingFuture() and it's inheritors

Towards #6197
diff --git a/api_core/google/api_core/future/polling.py b/api_core/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/api_core/google/api_core/operation.py b/api_core/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/api_core/tests/unit/test_operation.py b/api_core/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 = [
