chore(samples): harden 'tableadmin' samples against 429/503 errors (#418

) Closes #417.
googleapis · Oct 20, 2021 · 2841442 · 2841442
1 parent 4d392b2
commit 2841442
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/samples/tableadmin/requirements-test.txt b/samples/tableadmin/requirements-test.txt
@@ -1 +1,2 @@
 pytest==6.2.4
+google-cloud-testutils==1.0.0
diff --git a/samples/tableadmin/tableadmin_test.py b/samples/tableadmin/tableadmin_test.py
@@ -16,6 +16,9 @@
 import os
 import uuid
 
+from google.api_core import exceptions
+from test_utils.retry import RetryErrors
+
 from tableadmin import create_table
 from tableadmin import delete_table
 from tableadmin import run_table_operations
@@ -24,11 +27,13 @@
 BIGTABLE_INSTANCE = os.environ['BIGTABLE_INSTANCE']
 TABLE_ID_FORMAT = 'tableadmin-test-{}'
 
+retry_429_503 = RetryErrors(exceptions.TooManyRequests, exceptions.ServiceUnavailable)
+
 
 def test_run_table_operations(capsys):
     table_id = TABLE_ID_FORMAT.format(uuid.uuid4().hex[:8])
 
-    run_table_operations(PROJECT, BIGTABLE_INSTANCE, table_id)
+    retry_429_503(run_table_operations)(PROJECT, BIGTABLE_INSTANCE, table_id)
     out, _ = capsys.readouterr()
 
     assert 'Creating the ' + table_id + ' table.' in out
@@ -48,14 +53,14 @@ def test_run_table_operations(capsys):
     assert 'Delete a column family cf2...' in out
     assert 'Column family cf2 deleted successfully.' in out
 
-    delete_table(PROJECT, BIGTABLE_INSTANCE, table_id)
+    retry_429_503(delete_table)(PROJECT, BIGTABLE_INSTANCE, table_id)
 
 
 def test_delete_table(capsys):
     table_id = TABLE_ID_FORMAT.format(uuid.uuid4().hex[:8])
-    create_table(PROJECT, BIGTABLE_INSTANCE, table_id)
+    retry_429_503(create_table)(PROJECT, BIGTABLE_INSTANCE, table_id)
 
-    delete_table(PROJECT, BIGTABLE_INSTANCE, table_id)
+    retry_429_503(delete_table)(PROJECT, BIGTABLE_INSTANCE, table_id)
     out, _ = capsys.readouterr()
 
     assert 'Table ' + table_id + ' exists.' in out