diff options
author | Gurov Ilya <ilya.faer@mail.ru> | 2020-01-24 21:26:24 +0300 |
---|---|---|
committer | Bu Sun Kim <8822365+busunkim96@users.noreply.github.com> | 2020-01-24 10:26:24 -0800 |
commit | 14f1f34e013c90fed2da2918625083d299fda557 (patch) | |
tree | 629247fdcfcdab7cbd21bdc91794808679230a77 | |
parent | 18375fb453b0250ecaf7442c144ebe3b7b5e71ca (diff) | |
download | python-api-core-14f1f34e013c90fed2da2918625083d299fda557.tar.gz |
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
-rw-r--r-- | google/api_core/future/polling.py | 5 | ||||
-rw-r--r-- | google/api_core/operation.py | 17 | ||||
-rw-r--r-- | tests/unit/test_operation.py | 19 |
3 files changed, 35 insertions, 6 deletions
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 @@ class PollingFuture(base.Future): 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 @@ class Operation(polling.PollingFuture): ) 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 @@ 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 = [ |