test(data-science): add error checking, change machine type (GoogleCloudPlatform#4535)

Author: leahecole

data-science-onramp/data-cleaning/clean_test.py

@@ -2,9 +2,11 @@
 import re
 import uuid
 
+from google.api_core.exceptions import NotFound
 from google.cloud import bigquery
 from google.cloud import dataproc_v1 as dataproc
 from google.cloud import storage
+
 import pandas as pd
 import pytest
 
@@ -33,8 +35,12 @@
             "zone_uri": "",
             "metadata": {"PIP_PACKAGES": "google-cloud-storage"},
         },
-        "master_config": {"num_instances": 1, "machine_type_uri": "n1-standard-8"},
-        "worker_config": {"num_instances": 6, "machine_type_uri": "n1-standard-8"},
+        # We recommend these settings for running our code
+        # We use a less robust machine type for testing purposes
+        # "master_config": {"num_instances": 1, "machine_type_uri": "n1-standard-8"},
+        # "worker_config": {"num_instances": 6, "machine_type_uri": "n1-standard-8"},
+        "master_config": {"num_instances": 1, "machine_type_uri": "n1-standard-4"},
+        "worker_config": {"num_instances": 2, "machine_type_uri": "n1-standard-4"},
         "software_config": {
             "image_version": CLUSTER_IMAGE,
             # Change optional component to "ANACONDA" when this issue is resolved:
@@ -77,7 +83,10 @@ def setup_and_teardown_table():
     yield
 
     # Delete dataset
-    bq_client.delete_dataset(BQ_DATASET, delete_contents=True)
+    try:
+        bq_client.delete_dataset(BQ_DATASET, delete_contents=True)
+    except NotFound as e:
+        print(f"Ignoring NotFound upon cleanup, details: {e}")
 
 
 @pytest.fixture(autouse=True)

data-science-onramp/data-ingestion/setup_test.py

@@ -7,6 +7,7 @@
 import re
 import uuid
 
+from google.api_core.exceptions import NotFound
 from google.cloud import bigquery
 from google.cloud import dataproc_v1 as dataproc
 from google.cloud import storage
@@ -36,8 +37,12 @@
     "cluster_name": DATAPROC_CLUSTER,
     "config": {
         "gce_cluster_config": {"zone_uri": ""},
-        "master_config": {"num_instances": 1, "machine_type_uri": "n1-standard-8"},
-        "worker_config": {"num_instances": 6, "machine_type_uri": "n1-standard-8"},
+        # We recommend these configs when running the full code
+        # We use a less robust machine type for the tests
+        # "master_config": {"num_instances": 1, "machine_type_uri": "n1-standard-8"},
+        # "worker_config": {"num_instances": 6, "machine_type_uri": "n1-standard-8"},
+        "master_config": {"num_instances": 1, "machine_type_uri": "n1-standard-4"},
+        "worker_config": {"num_instances": 2, "machine_type_uri": "n1-standard-4"},
         "software_config": {
             "image_version": CLUSTER_IMAGE,
             "optional_components": [5],
@@ -102,7 +107,10 @@ def setup_and_teardown_bq_dataset():
     yield
 
     # Delete Dataset
-    bq_client.delete_dataset(BQ_DATASET, delete_contents=True)
+    try:
+        bq_client.delete_dataset(BQ_DATASET, delete_contents=True)
+    except NotFound as e:
+        print(f"Ignoring NotFound on cleanup, details: {e}")
 
 
 def get_blob_from_path(path):