From 284144294a715d596aab71ef042bee121c80f6e1 Mon Sep 17 00:00:00 2001
From: Tres Seaver <tseaver@palladion.com>
Date: Wed, 20 Oct 2021 14:13:53 -0400
Subject: [PATCH] chore(samples): harden 'tableadmin' samples against 429/503
 errors (#418)

Closes #417.
---
 samples/tableadmin/requirements-test.txt |  1 +
 samples/tableadmin/tableadmin_test.py    | 13 +++++++++----
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/samples/tableadmin/requirements-test.txt b/samples/tableadmin/requirements-test.txt
index 95ea1e6a0..2ff95fe08 100644
--- 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
index c0ef09d12..b001ce076 100755
--- 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