test(spanner): fix flaky restore_database() sample test (#4641)

* test(spanner): fix flaky restore backup sample test

* fix: update timeout to match other languages

* test: import test_utils and set max retries to 2

Co-authored-by: larkee <[email protected]>

Author: larkee

spanner/cloud-client/backup_sample.py

@@ -68,7 +68,7 @@ def restore_database(instance_id, new_database_id, backup_id):
     operation = new_database.restore(backup)
 
     # Wait for restore operation to complete.
-    operation.result(1200)
+    operation.result(1600)
 
     # Newly created database has restore information.
     new_database.reload()

spanner/cloud-client/backup_sample_test.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 import uuid
 
+from google.api_core.exceptions import DeadlineExceeded
 from google.cloud import spanner
 import pytest
+from test_utils.retry import RetryErrors
 
 import backup_sample
 
@@ -67,6 +69,7 @@ def test_create_backup(capsys, database):
     assert BACKUP_ID in out
 
 
+@RetryErrors(exception=DeadlineExceeded, max_tries=2)
 def test_restore_database(capsys):
     backup_sample.restore_database(INSTANCE_ID, RESTORE_DB_ID, BACKUP_ID)
     out, _ = capsys.readouterr()

spanner/cloud-client/requirements-test.txt

@@ -1,2 +1,3 @@
 pytest==6.0.1
 mock==4.0.2
+google-cloud-testutils==0.1.0