diff --git a/dataproc/snippets/submit_pyspark_job_to_driver_node_group_cluster_test.py b/dataproc/snippets/submit_pyspark_job_to_driver_node_group_cluster_test.py index 38e3ebb24e3..9b156e7147b 100644 --- a/dataproc/snippets/submit_pyspark_job_to_driver_node_group_cluster_test.py +++ b/dataproc/snippets/submit_pyspark_job_to_driver_node_group_cluster_test.py @@ -62,27 +62,28 @@ def teardown(): max_tries=5, ) def test_workflows(capsys): - # Setup driver node group cluster. TODO: cleanup b/424371877 - command = f"""gcloud dataproc clusters create {CLUSTER_NAME} \ - --region {REGION} \ - --project {PROJECT_ID} \ - --driver-pool-size=1 \ - --driver-pool-id=pytest""" - - output = subprocess.run( - command, - capture_output=True, - shell=True, - check=True, - ) - print(output) + try: + # Setup driver node group cluster. TODO: cleanup b/424371877 + command = f"""gcloud dataproc clusters create {CLUSTER_NAME} \ + --region {REGION} \ + --project {PROJECT_ID} \ + --driver-pool-size=1 \ + --driver-pool-id=pytest""" - # Wrapper function for client library function - submit_pyspark_job_to_driver_node_group_cluster.submit_job( - PROJECT_ID, REGION, CLUSTER_NAME - ) + output = subprocess.run( + command, + capture_output=True, + shell=True, + check=True, + ) + print(output) - out, _ = capsys.readouterr() - assert "Job finished successfully" in out + # Wrapper function for client library function + submit_pyspark_job_to_driver_node_group_cluster.submit_job( + PROJECT_ID, REGION, CLUSTER_NAME + ) - # cluster deleted in teardown() + out, _ = capsys.readouterr() + assert "Job finished successfully" in out + finally: + teardown() diff --git a/dataproc/snippets/submit_spark_job_to_driver_node_group_cluster_test.py b/dataproc/snippets/submit_spark_job_to_driver_node_group_cluster_test.py index ac642ed2e5a..fa7d819b6cf 100644 --- a/dataproc/snippets/submit_spark_job_to_driver_node_group_cluster_test.py +++ b/dataproc/snippets/submit_spark_job_to_driver_node_group_cluster_test.py @@ -62,27 +62,28 @@ def teardown(): max_tries=5, ) def test_workflows(capsys): - # Setup driver node group cluster. TODO: cleanup b/424371877 - command = f"""gcloud dataproc clusters create {CLUSTER_NAME} \ - --region {REGION} \ - --project {PROJECT_ID} \ - --driver-pool-size=1 \ - --driver-pool-id=pytest""" - - output = subprocess.run( - command, - capture_output=True, - shell=True, - check=True, - ) - print(output) + try: + # Setup driver node group cluster. TODO: cleanup b/424371877 + command = f"""gcloud dataproc clusters create {CLUSTER_NAME} \ + --region {REGION} \ + --project {PROJECT_ID} \ + --driver-pool-size=1 \ + --driver-pool-id=pytest""" - # Wrapper function for client library function - submit_spark_job_to_driver_node_group_cluster.submit_job( - PROJECT_ID, REGION, CLUSTER_NAME - ) + output = subprocess.run( + command, + capture_output=True, + shell=True, + check=True, + ) + print(output) - out, _ = capsys.readouterr() - assert "Job finished successfully" in out + # Wrapper function for client library function + submit_spark_job_to_driver_node_group_cluster.submit_job( + PROJECT_ID, REGION, CLUSTER_NAME + ) - # cluster deleted in teardown() + out, _ = capsys.readouterr() + assert "Job finished successfully" in out + finally: + teardown()