diff --git a/dataproc/snippets/create_cluster_test.py b/dataproc/snippets/create_cluster_test.py index c180c4abd5..fb0d4bb7e9 100644 --- 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 index a3eef34047..e9395659c7 100644 --- 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 index 3b35e3e5b3..df421e6707 100644 --- 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)