add more error handling

Author: leahecole

dataproc/snippets/create_cluster_test.py

@@ -12,11 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import datetime
 import os
 import uuid
 
 import backoff
-from google.api_core.exceptions import (InternalServerError, InvalidArgument, NotFound,
+from google.api_core.exceptions import (AlreadyExists, InternalServerError, InvalidArgument, NotFound,
                                         ServiceUnavailable)
 from google.cloud import dataproc_v1 as dataproc
 
@@ -26,12 +27,21 @@
 REGION = "us-west1"
 CLUSTER_NAME = "py-cc-test-{}".format(str(uuid.uuid4()))
 
+cluster_client = dataproc.ClusterControllerClient(
+    client_options={"api_endpoint": f"{REGION}-dataproc.googleapis.com:443"}
+)
+
 
 @backoff.on_exception(backoff.expo, (Exception), max_tries=5)
 def teardown():
+<<<<<<< HEAD
     cluster_client = dataproc.ClusterControllerClient(
         client_options={"api_endpoint": f"{REGION}-dataproc.googleapis.com:443"}
     )
+=======
+    yield
+
+>>>>>>> ceedc468a (add more error handling)
     # Client library function
     try:
         operation = cluster_client.delete_cluster(
@@ -53,8 +63,21 @@ def test_cluster_create(capsys):
     # Wrapper function for client library function
     try:
         create_cluster.create_cluster(PROJECT_ID, REGION, CLUSTER_NAME)
-    finally:
-        teardown()
+        out, _ = capsys.readouterr()
+        assert CLUSTER_NAME in out
+    except AlreadyExists:
+        request = dataproc.GetClusterRequest(project_id=PROJECT_ID, region=REGION, cluster_name=CLUSTER_NAME)
+        response = cluster_client.get_cluster(request=request)
+        assert response.status.state == dataproc.ClusterStatus.State(2)  # verify the cluster is in the RUNNING state
 
-    out, _ = capsys.readouterr()
-    assert CLUSTER_NAME in out
+        status_start = response.status.state_start_time  # when the cluster started being in the RUNNING state
+        now = datetime.datetime.now(datetime.timezone.utc)
+        diff = now-status_start
+        # check that it's been running for less than 20 min
+        # this means we probably did a backoff during the creation for a different reason
+        # and that this cluster was created as part of this test, so we can continue
+        assert diff.seconds/60 < 20
+        out, _ = capsys.readouterr()
+        assert CLUSTER_NAME in out
+    finally:
+        teardown()

dataproc/snippets/submit_job_test.py

@@ -16,7 +16,7 @@
 import uuid
 
 import backoff
-from google.api_core.exceptions import (AlreadyExists, InternalServerError, NotFound,
+from google.api_core.exceptions import (AlreadyExists, InternalServerError, InvalidArgument, NotFound,
                                         ServiceUnavailable)
 from google.cloud import dataproc_v1 as dataproc
 import pytest
@@ -85,7 +85,8 @@ def cluster_name(cluster_client):
         teardown_cluster(cluster_client, curr_cluster_name)
 
 
-@backoff.on_exception(backoff.expo, (InternalServerError, ServiceUnavailable), max_tries=5)
+# InvalidArgument is thrown when the subnetwork is not ready
+@backoff.on_exception(backoff.expo, (InvalidArgument, InternalServerError, ServiceUnavailable), max_tries=5)
 def test_submit_job(capsys, cluster_name):
     submit_job.submit_job(PROJECT_ID, REGION, cluster_name)
     out, _ = capsys.readouterr()

dataproc/snippets/update_cluster_test.py

@@ -20,6 +20,7 @@
 
 import backoff
 from google.api_core.exceptions import (
+    AlreadyExists,
     Cancelled,
     InternalServerError,
     InvalidArgument,
@@ -56,7 +57,7 @@ def cluster_client():
 
 
 @backoff.on_exception(backoff.expo, (ServiceUnavailable, InvalidArgument), max_tries=5)
-def setup_cluster(cluster_client):
+def setup(cluster_client):
     # Create the cluster.
     operation = cluster_client.create_cluster(
         request={"project_id": PROJECT_ID, "region": REGION, "cluster": CLUSTER}
@@ -75,8 +76,28 @@ def teardown_cluster(cluster_client):
             }
         )
         operation.result()
-    except NotFound:
-        print("Cluster already deleted")
+
+    def teardown():
+        try:
+            operation = cluster_client.delete_cluster(
+                request={
+                    "project_id": PROJECT_ID,
+                    "region": REGION,
+                    "cluster_name": CLUSTER_NAME,
+                }
+            )
+            operation.result()
+        except NotFound:
+            print("Cluster already deleted")
+
+    try:
+        setup()
+        yield
+    except AlreadyExists:
+        print("Cluster exists, utilizing existing cluster")
+        yield
+    finally:
+        teardown()
 
 
 @backoff.on_exception(