test: add retries to dataproc tests (#8874)

* test: add retries to dataproc tests * reorder imports * linting Co-authored-by: nicain <[email protected]>
GoogleCloudPlatform · Mar 13, 2023 · 0013c10 · 0013c10
1 parent 18b889f
commit 0013c10
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 11 deletions.
diff --git a/dataproc/snippets/create_cluster_test.py b/dataproc/snippets/create_cluster_test.py
@@ -15,13 +15,14 @@
 import os
 import uuid
 
-from google.api_core.exceptions import NotFound
+import backoff
+from google.api_core.exceptions import (InternalServerError, NotFound,
+                                        ServiceUnavailable)
 from google.cloud import dataproc_v1 as dataproc
 import pytest
 
 import create_cluster
 
-
 PROJECT_ID = os.environ["GOOGLE_CLOUD_PROJECT"]
 REGION = "us-central1"
 CLUSTER_NAME = "py-cc-test-{}".format(str(uuid.uuid4()))
@@ -49,6 +50,7 @@ def teardown():
         print("Cluster already deleted")
 
 
+@backoff.on_exception(backoff.expo, (InternalServerError, ServiceUnavailable), max_tries=5)
 def test_cluster_create(capsys):
     # Wrapper function for client library function
     create_cluster.create_cluster(PROJECT_ID, REGION, CLUSTER_NAME)

diff --git a/dataproc/snippets/instantiate_inline_workflow_template_test.py b/dataproc/snippets/instantiate_inline_workflow_template_test.py
@@ -14,13 +14,16 @@
 
 import os
 
-import instantiate_inline_workflow_template
+import backoff
+from google.api_core.exceptions import InternalServerError, ServiceUnavailable
 
+import instantiate_inline_workflow_template
 
 PROJECT_ID = os.environ["GOOGLE_CLOUD_PROJECT"]
 REGION = "us-central1"
 
 
+@backoff.on_exception(backoff.expo, (InternalServerError, ServiceUnavailable), max_tries=5)
 def test_workflows(capsys):
     # Wrapper function for client library function
     instantiate_inline_workflow_template.instantiate_inline_workflow_template(

diff --git a/dataproc/snippets/submit_job_test.py b/dataproc/snippets/submit_job_test.py
@@ -15,13 +15,14 @@
 import os
 import uuid
 
-from google.api_core.exceptions import NotFound
+import backoff
+from google.api_core.exceptions import (InternalServerError, NotFound,
+                                        ServiceUnavailable)
 from google.cloud import dataproc_v1 as dataproc
 import pytest
 
 import submit_job
 
-
 PROJECT_ID = os.environ["GOOGLE_CLOUD_PROJECT"]
 REGION = "us-central1"
 CLUSTER_NAME = "py-sj-test-{}".format(str(uuid.uuid4()))
@@ -67,6 +68,7 @@ def setup_teardown():
             print("Cluster already deleted")
 
 
+@backoff.on_exception(backoff.expo, (InternalServerError, ServiceUnavailable), max_tries=5)
 def test_submit_job(capsys):
     submit_job.submit_job(PROJECT_ID, REGION, CLUSTER_NAME)
     out, _ = capsys.readouterr()

diff --git a/dataproc/snippets/update_cluster_test.py b/dataproc/snippets/update_cluster_test.py
@@ -15,19 +15,18 @@
 # This sample walks a user through updating the number of clusters using the Dataproc
 # client library.
 
-
 import os
 import uuid
 
-from google.api_core.exceptions import NotFound
-from google.cloud.dataproc_v1.services.cluster_controller.client import (
-    ClusterControllerClient,
-)
+import backoff
+from google.api_core.exceptions import (InternalServerError, NotFound,
+                                        ServiceUnavailable)
+from google.cloud.dataproc_v1.services.cluster_controller.client import \
+    ClusterControllerClient
 import pytest
 
 import update_cluster
 
-
 PROJECT_ID = os.environ["GOOGLE_CLOUD_PROJECT"]
 REGION = "us-central1"
 CLUSTER_NAME = f"py-cc-test-{str(uuid.uuid4())}"
@@ -74,6 +73,7 @@ def setup_teardown(cluster_client):
             print("Cluster already deleted")
 
 
+@backoff.on_exception(backoff.expo, (InternalServerError, ServiceUnavailable), max_tries=5)
 def test_update_cluster(capsys, cluster_client: ClusterControllerClient):
     # Wrapper function for client library function
     update_cluster.update_cluster(PROJECT_ID, REGION, CLUSTER_NAME, NEW_NUM_INSTANCES)