fix(sample-test): add backoff around the cleanup code (#65)

googleapis · Oct 16, 2020 · 5b22fa8 · 5b22fa8
1 parent 14571c7
commit 5b22fa8
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/packages/google-cloud-dlp/samples/snippets/inspect_content_test.py b/packages/google-cloud-dlp/samples/snippets/inspect_content_test.py
@@ -15,7 +15,9 @@
 import os
 import uuid
 
+import backoff
 import google.api_core.exceptions
+from google.api_core.exceptions import ServiceUnavailable
 import google.cloud.bigquery
 import google.cloud.datastore
 import google.cloud.dlp_v2
@@ -122,7 +124,11 @@ def datastore_project():
 
     yield GCLOUD_PROJECT
 
-    datastore_client.delete(key)
+    @backoff.on_exception(backoff.expo, ServiceUnavailable, max_time=120)
+    def cleanup():
+        datastore_client.delete(key)
+
+    cleanup()
 
 
 @pytest.fixture(scope="module")
@@ -157,7 +163,11 @@ def bigquery_project():
 
     yield GCLOUD_PROJECT
 
-    bigquery_client.delete_dataset(dataset_ref, delete_contents=True)
+    @backoff.on_exception(backoff.expo, ServiceUnavailable, max_time=120)
+    def cleanup():
+        bigquery_client.delete_dataset(dataset_ref, delete_contents=True)
+
+    cleanup()
 
 
 def test_inspect_string_basic(capsys):

diff --git a/packages/google-cloud-dlp/samples/snippets/requirements-test.txt b/packages/google-cloud-dlp/samples/snippets/requirements-test.txt
@@ -1,4 +1,4 @@
+backoff==1.10.0
 pytest==6.0.1
 flaky==3.7.0
 mock==4.0.2
-