fix: changes default region for dataproc quickstart test to us-central2 (#9764)

* fix: changes default region for dataproc quickstart test to us-central2

* Update dataproc/snippets/quickstart/quickstart_test.py

Co-authored-by: Leah E. Cole <[email protected]>

* changing to regions with unlimited persistent disk quota

* refactor test_submit_job to utilize backoff retries correctly, and set disk_config to smaller values for several tests

* reverting to us-central1

---------

Co-authored-by: Leah E. Cole <[email protected]>

Author: nicain

dataproc/snippets/submit_job_test.py

@@ -16,7 +16,7 @@
 import uuid
 
 import backoff
-from google.api_core.exceptions import (InternalServerError, InvalidArgument, NotFound,
+from google.api_core.exceptions import (InternalServerError, NotFound,
                                         ServiceUnavailable)
 from google.cloud import dataproc_v1 as dataproc
 import pytest
@@ -25,57 +25,66 @@
 
 PROJECT_ID = os.environ["GOOGLE_CLOUD_PROJECT"]
 REGION = "us-central1"
-CLUSTER_NAME = "py-sj-test-{}".format(str(uuid.uuid4()))
-CLUSTER = {
-    "project_id": PROJECT_ID,
-    "cluster_name": CLUSTER_NAME,
-    "config": {
-        "master_config": {"num_instances": 1, "machine_type_uri": "n1-standard-2"},
-        "worker_config": {"num_instances": 2, "machine_type_uri": "n1-standard-2"},
-    },
-}
-
-
-@pytest.fixture(autouse=True)
-def setup_teardown():
-    cluster_client = dataproc.ClusterControllerClient(
+
+
+@pytest.fixture(scope='module')
+def cluster_client():
+    return 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}
+
+@backoff.on_exception(backoff.expo, ServiceUnavailable, max_tries=5)
+def setup_cluster(cluster_client, curr_cluster_name):
+
+    CLUSTER = {
+        "project_id": PROJECT_ID,
+        "cluster_name": curr_cluster_name,
+        "config": {
+            "master_config": {"num_instances": 1, "machine_type_uri": "n1-standard-2", "disk_config": {"boot_disk_size_gb": 100}},
+            "worker_config": {"num_instances": 2, "machine_type_uri": "n1-standard-2", "disk_config": {"boot_disk_size_gb": 100}},
+        },
+    }
+
+    # Create the cluster.
+    operation = cluster_client.create_cluster(
+        request={"project_id": PROJECT_ID, "region": REGION, "cluster": CLUSTER}
+    )
+    operation.result()
+
+
+@backoff.on_exception(backoff.expo, ServiceUnavailable, max_tries=5)
+def teardown_cluster(cluster_client, curr_cluster_name):
+    try:
+        operation = cluster_client.delete_cluster(
+            request={
+                "project_id": PROJECT_ID,
+                "region": REGION,
+                "cluster_name": curr_cluster_name,
+            }
         )
         operation.result()
 
-    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")
+    except NotFound:
+        print("Cluster already deleted")
+
+
+@pytest.fixture(scope='module')
+def cluster_name(cluster_client):
+    curr_cluster_name = "py-sj-test-{}".format(str(uuid.uuid4()))
+
     try:
-        setup()
-        yield
+        setup_cluster(cluster_client, curr_cluster_name)
+        yield curr_cluster_name
     finally:
-        teardown()
+        teardown_cluster(cluster_client, curr_cluster_name)
 
 
 @backoff.on_exception(backoff.expo, (InternalServerError, ServiceUnavailable), max_tries=5)
-def test_submit_job(capsys):
-    submit_job.submit_job(PROJECT_ID, REGION, CLUSTER_NAME)
+def test_submit_job(capsys, cluster_name):
+    submit_job.submit_job(PROJECT_ID, REGION, cluster_name)
     out, _ = capsys.readouterr()
 
     assert "Job finished successfully" in out

dataproc/snippets/update_cluster_test.py

@@ -41,8 +41,8 @@
     "project_id": PROJECT_ID,
     "cluster_name": CLUSTER_NAME,
     "config": {
-        "master_config": {"num_instances": 1, "machine_type_uri": "n1-standard-2"},
-        "worker_config": {"num_instances": 2, "machine_type_uri": "n1-standard-2"},
+        "master_config": {"num_instances": 1, "machine_type_uri": "n1-standard-2", "disk_config": {"boot_disk_size_gb": 100}},
+        "worker_config": {"num_instances": 2, "machine_type_uri": "n1-standard-2", "disk_config": {"boot_disk_size_gb": 100}},
     },
 }
 

pubsublite/spark-connector/spark_streaming_test.py

@@ -124,8 +124,8 @@ def dataproc_cluster() -> Generator[dataproc_v1.Cluster, None, None]:
         "project_id": PROJECT_ID,
         "cluster_name": CLUSTER_ID,
         "config": {
-            "master_config": {"num_instances": 1, "machine_type_uri": "n1-standard-2"},
-            "worker_config": {"num_instances": 2, "machine_type_uri": "n1-standard-2"},
+            "master_config": {"num_instances": 1, "machine_type_uri": "n1-standard-2", "disk_config": {"boot_disk_size_gb": 100}},
+            "worker_config": {"num_instances": 2, "machine_type_uri": "n1-standard-2", "disk_config": {"boot_disk_size_gb": 100}},
             "config_bucket": BUCKET,
             "temp_bucket": BUCKET,
             "software_config": {"image_version": "2.0-debian10"},