add backoff to fix flaky test (GoogleCloudPlatform#8054)

leahecole · web-flow · commit d5d980af8a31 · 2022-06-07T12:45:16.000-04:00
* add backoff to fix flaky test

* fix lint
diff --git a/composer/2022_airflow_summit/data_analytics_process_test.py b/composer/2022_airflow_summit/data_analytics_process_test.py
@@ -20,7 +20,8 @@
 import os
 import uuid
 
-from google.api_core.exceptions import NotFound
+import backoff
+from google.api_core.exceptions import Aborted, NotFound
 from google.cloud import bigquery
 from google.cloud import dataproc_v1 as dataproc
 from google.cloud import storage
@@ -122,13 +123,13 @@ 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_bucket):
     # check that the results table isnt there
     with pytest.raises(NotFound):
         BQ_CLIENT.get_table(f"{BQ_DATASET}.{BQ_WRITE_TABLE}")
 
-    # do the process
-    # assert the column is there
     # create a batch
     dataproc_client = dataproc.BatchControllerClient(
         client_options={
diff --git a/composer/2022_airflow_summit/requirements-test.txt b/composer/2022_airflow_summit/requirements-test.txt
@@ -1,3 +1,4 @@
 pytest==7.0.1
 cloud-composer-dag-test-utils==1.0.0
 markupsafe==2.0.1
+backoff==2.0.1
