add retries for subnetwork not ready errors (GoogleCloudPlatform#9523)

leahecole · web-flow · commit abeeaff58610 · 2023-04-10T12:02:56.000-04:00
* add retries for subnetwork not ready errors

* fix lint

* experiment: remove autouse

* chain decorators

* fix lint

* refactor fixtures with inner functions
diff --git a/dataproc/snippets/create_cluster_test.py b/dataproc/snippets/create_cluster_test.py
@@ -16,7 +16,7 @@
 import uuid
 
 import backoff
-from google.api_core.exceptions import (InternalServerError, NotFound,
+from google.api_core.exceptions import (InternalServerError, InvalidArgument, NotFound,
                                         ServiceUnavailable)
 from google.cloud import dataproc_v1 as dataproc
 import pytest
@@ -50,7 +50,8 @@ def teardown():
         print("Cluster already deleted")
 
 
-@backoff.on_exception(backoff.expo, (InternalServerError, ServiceUnavailable), max_tries=5)
+# InvalidArgument is thrown when the subnetwork is not ready
+@backoff.on_exception(backoff.expo, (InternalServerError, ServiceUnavailable, InvalidArgument), 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/submit_job_test.py b/dataproc/snippets/submit_job_test.py
@@ -16,7 +16,7 @@
 import uuid
 
 import backoff
-from google.api_core.exceptions import (InternalServerError, NotFound,
+from google.api_core.exceptions import (InternalServerError, InvalidArgument, NotFound,
                                         ServiceUnavailable)
 from google.cloud import dataproc_v1 as dataproc
 import pytest
@@ -38,22 +38,22 @@
 
 @pytest.fixture(autouse=True)
 def setup_teardown():
-    try:
-        cluster_client = dataproc.ClusterControllerClient(
-            client_options={
-                "api_endpoint": "{}-dataproc.googleapis.com:443".format(REGION)
-            }
-        )
+    cluster_client = dataproc.ClusterControllerClient(
+        client_options={
+            "api_endpoint": "{}-dataproc.googleapis.com:443".format(REGION)
+        }
+    )
 
+    # Retry on InvalidArgument subnetwork not ready error
+    @backoff.on_exception(backoff.expo, (InvalidArgument), max_tries=3)
+    def setup():
         # Create the cluster.
         operation = cluster_client.create_cluster(
             request={"project_id": PROJECT_ID, "region": REGION, "cluster": CLUSTER}
         )
         operation.result()
 
-        yield
-
-    finally:
+    def teardown():
         try:
             operation = cluster_client.delete_cluster(
                 request={
@@ -66,6 +66,11 @@ def setup_teardown():
 
         except NotFound:
             print("Cluster already deleted")
+    try:
+        setup()
+        yield
+    finally:
+        teardown()
 
 
 @backoff.on_exception(backoff.expo, (InternalServerError, ServiceUnavailable), max_tries=5)
diff --git a/dataproc/snippets/update_cluster_test.py b/dataproc/snippets/update_cluster_test.py
@@ -19,7 +19,7 @@
 import uuid
 
 import backoff
-from google.api_core.exceptions import (InternalServerError, NotFound,
+from google.api_core.exceptions import (InternalServerError, InvalidArgument, NotFound,
                                         ServiceUnavailable)
 from google.cloud.dataproc_v1.services.cluster_controller.client import \
     ClusterControllerClient
@@ -51,15 +51,16 @@ def cluster_client():
 
 @pytest.fixture(autouse=True)
 def setup_teardown(cluster_client):
-    try:
+    # InvalidArgument is thrown when the subnetwork is not ready
+    @backoff.on_exception(backoff.expo, (InvalidArgument), max_tries=3)
+    def setup():
         # Create the cluster.
         operation = cluster_client.create_cluster(
             request={"project_id": PROJECT_ID, "region": REGION, "cluster": CLUSTER}
         )
         operation.result()
 
-        yield
-    finally:
+    def teardown():
         try:
             operation = cluster_client.delete_cluster(
                 request={
@@ -71,6 +72,11 @@ def setup_teardown(cluster_client):
             operation.result()
         except NotFound:
             print("Cluster already deleted")
+    try:
+        setup()
+        yield
+    finally:
+        teardown()
 
 
 @backoff.on_exception(backoff.expo, (InternalServerError, ServiceUnavailable), max_tries=5)
