diff --git a/logging/tests/system/test_system.py b/logging/tests/system/test_system.py index 5dc8fd670494..ccfac6d7b272 100644 --- a/logging/tests/system/test_system.py +++ b/logging/tests/system/test_system.py @@ -16,6 +16,7 @@ import logging import unittest +from google.api_core.exceptions import BadGateway from google.api_core.exceptions import Conflict from google.api_core.exceptions import NotFound from google.api_core.exceptions import TooManyRequests @@ -102,7 +103,7 @@ def setUp(self): self._handlers_cache = logging.getLogger().handlers[:] def tearDown(self): - retry = RetryErrors(NotFound, max_tries=9) + retry = RetryErrors((NotFound, TooManyRequests), max_tries=9) for doomed in self.to_delete: try: retry(doomed.delete)() @@ -381,9 +382,10 @@ def _init_storage_bucket(self): # Create the destination bucket, and set up the ACL to allow # Stackdriver Logging to write into it. + retry = RetryErrors((Conflict, TooManyRequests, ServiceUnavailable)) storage_client = storage.Client() bucket = storage_client.bucket(BUCKET_NAME) - retry_429(bucket.create)() + retry(bucket.create)() self.to_delete.append(bucket) bucket.acl.reload() logs_group = bucket.acl.group('cloud-logs@google.com') @@ -441,9 +443,11 @@ def _init_bigquery_dataset(self): # Create the destination dataset, and set up the ACL to allow # Stackdriver Logging to write into it. + retry = RetryErrors((TooManyRequests, BadGateway, ServiceUnavailable)) bigquery_client = bigquery.Client() dataset_ref = bigquery_client.dataset(dataset_name) - dataset = bigquery_client.create_dataset(bigquery.Dataset(dataset_ref)) + dataset = retry(bigquery_client.create_dataset)( + bigquery.Dataset(dataset_ref)) self.to_delete.append((bigquery_client, dataset)) bigquery_client.get_dataset(dataset) access = AccessEntry(