refactor test to match #8328

Author: leahecole

composer/2022_airflow_summit/data_analytics_process_expansion_test.py

@@ -20,7 +20,6 @@
 import os
 import uuid
 
-import backoff
 from google.api_core.exceptions import Aborted, NotFound
 from google.cloud import bigquery
 from google.cloud import dataproc_v1 as dataproc
@@ -58,8 +57,22 @@
 PROCESSING_PYTHON_FILE = f"gs://{BUCKET_NAME}/{BUCKET_BLOB}"
 
 
-@pytest.fixture(scope="module")
-def test_dataproc_batch():
+@pytest.fixture(scope="function")
+def test_dataproc_batch(test_bucket, bq_dataset):
+    # check that the results tables aren't there
+    # considered using pytest parametrize, but did not want rest of test
+    # to run 5 times - only this part
+    output_tables = [
+        BQ_NORMALIZED_TABLE_NAME,
+        BQ_PRCP_MEAN_TABLE_NAME,
+        BQ_SNOW_MEAN_TABLE_NAME,
+        BQ_PHX_PRCP_TABLE_NAME,
+        BQ_PHX_SNOW_TABLE_NAME,
+    ]
+    for output_table in output_tables:
+        with pytest.raises(NotFound):
+            BQ_CLIENT.get_table(f"{BQ_DESTINATION_DATASET_NAME}.{output_table}")
+
 
     BATCH_ID = f"summit-dag-expansion-test-{TEST_ID}"  # Dataproc serverless only allows lowercase characters
     BATCH_CONFIG = {
@@ -78,7 +91,47 @@ def test_dataproc_batch():
         }
     }
 
-    yield (BATCH_ID, BATCH_CONFIG)
+    # create a batch
+    dataproc_client = dataproc.BatchControllerClient(
+        client_options={
+            "api_endpoint": f"{DATAPROC_REGION}-dataproc.googleapis.com:443"
+        }
+    )
+    request = dataproc.CreateBatchRequest(
+        parent=f"projects/{PROJECT_ID}/regions/{DATAPROC_REGION}",
+        batch=test_dataproc_batch[1],
+        batch_id=test_dataproc_batch[0],
+    )
+    try:
+        # Make the request
+        operation = dataproc_client.create_batch(request=request)
+
+        print("Waiting for operation to complete...")
+
+        response = operation.result()
+    except Aborted as e:
+        # retry once if we see a flaky 409 "subnet not ready error"
+        if "/subnetworks/default" in e:
+            # delete the errored out batch so we don't see an "AlreadyExists"
+            delete_request = dataproc.DeleteBatchRequest(
+                name=f"projects/{PROJECT_ID}/locations/{DATAPROC_REGION}/batches/{BATCH_ID}"
+            )
+            dataproc_client.delete_batch(request=delete_request)
+            # retry the creation operation once
+            create_request = dataproc.CreateBatchRequest(
+                parent=f"projects/{PROJECT_ID}/regions/{DATAPROC_REGION}",
+                batch=BATCH_CONFIG,
+                batch_id=BATCH_ID,
+            )
+            operation = dataproc_client.create_batch(request=create_request)
+
+            print("Waiting for operation to complete...")
+
+            response = operation.result()
+        else:
+            raise (e)
+
+    yield response
     dataproc_client = dataproc.BatchControllerClient(
         client_options={
             "api_endpoint": f"{DATAPROC_REGION}-dataproc.googleapis.com:443"
@@ -162,43 +215,10 @@ def bq_dataset(test_bucket):
         print(f"Ignoring NotFound on cleanup, details: {e}")
 
 
-# Retry if we see a flaky 409 "subnet not ready" exception
-@backoff.on_exception(backoff.expo, Aborted, max_tries=3)
-def test_process(test_dataproc_batch):
-    # check that the results tables aren't there
-    # considered using pytest parametrize, but did not want rest of test
-    # to run 5 times - only this part
-    output_tables = [
-        BQ_NORMALIZED_TABLE_NAME,
-        BQ_PRCP_MEAN_TABLE_NAME,
-        BQ_SNOW_MEAN_TABLE_NAME,
-        BQ_PHX_PRCP_TABLE_NAME,
-        BQ_PHX_SNOW_TABLE_NAME,
-    ]
-    for output_table in output_tables:
-        with pytest.raises(NotFound):
-            BQ_CLIENT.get_table(f"{BQ_DESTINATION_DATASET_NAME}.{output_table}")
-
-    # create a batch
-    dataproc_client = dataproc.BatchControllerClient(
-        client_options={
-            "api_endpoint": f"{DATAPROC_REGION}-dataproc.googleapis.com:443"
-        }
-    )
-    request = dataproc.CreateBatchRequest(
-        parent=f"projects/{PROJECT_ID}/regions/{DATAPROC_REGION}",
-        batch=test_dataproc_batch[1],
-        batch_id=test_dataproc_batch[0],
-    )
-    # Make the request
-    operation = dataproc_client.create_batch(request=request)
 
-    print("Waiting for operation to complete...")
-
-    response = operation.result()
-
-    # Handle the response
-    print(response)
+def test_process(test_dataproc_batch):
+    print(test_dataproc_batch)
+ 
 
     # check that the results table is there now
     assert (