migrate system test gcs_to_bigquery into new design (#22753)

joppevos · web-flow · commit 9a623e94cb3e · 2022-04-13T21:30:05.000+02:00
diff --git a/docs/apache-airflow-providers-google/operators/cloud/gcs.rst b/docs/apache-airflow-providers-google/operators/cloud/gcs.rst
@@ -41,7 +41,7 @@ Use the
 :class:`~airflow.providers.google.cloud.transfers.gcs_to_bigquery.GCSToBigQueryOperator`
 to execute a BigQuery load job.
 
-.. exampleinclude:: /../../airflow/providers/google/cloud/example_dags/example_gcs_to_bigquery.py
+.. exampleinclude:: /../../tests/system/providers/google/gcs/example_gcs_to_bigquery.py
     :language: python
     :dedent: 4
     :start-after: [START howto_operator_gcs_to_bigquery]
diff --git a/tests/providers/google/cloud/transfers/test_gcs_to_bigquery_system.py b/tests/providers/google/cloud/transfers/test_gcs_to_bigquery_system.py
diff --git a/tests/system/providers/google/gcs/example_gcs_to_bigquery.py b/tests/system/providers/google/gcs/example_gcs_to_bigquery.py
@@ -29,19 +29,24 @@
     BigQueryDeleteDatasetOperator,
 )
 from airflow.providers.google.cloud.transfers.gcs_to_bigquery import GCSToBigQueryOperator
+from airflow.utils.trigger_rule import TriggerRule
 
-DATASET_NAME = os.environ.get("GCP_DATASET_NAME", 'airflow_test')
-TABLE_NAME = os.environ.get("GCP_TABLE_NAME", 'gcs_to_bq_table')
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "gcs_to_bigquery_operator"
+
+DATASET_NAME = f"dataset_{DAG_ID}_{ENV_ID}"
+TABLE_NAME = "test"
+PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT")
 
 with models.DAG(
-    dag_id='example_gcs_to_bigquery_operator',
+    dag_id=DAG_ID,
+    schedule_interval='@once',
     start_date=datetime(2021, 1, 1),
     catchup=False,
-    schedule_interval='@once',
-    tags=['example'],
+    tags=['example', "gcs"],
 ) as dag:
     create_test_dataset = BigQueryCreateEmptyDatasetOperator(
-        task_id='create_airflow_test_dataset', dataset_id=DATASET_NAME
+        task_id='create_airflow_test_dataset', dataset_id=DATASET_NAME, project_id=PROJECT_ID
     )
 
     # [START howto_operator_gcs_to_bigquery]
@@ -62,6 +67,25 @@
         task_id='delete_airflow_test_dataset',
         dataset_id=DATASET_NAME,
         delete_contents=True,
+        trigger_rule=TriggerRule.ALL_DONE,
+    )
+
+    (
+        # TEST SETUP
+        create_test_dataset
+        # TEST BODY
+        >> load_csv
+        # TEST TEARDOWN
+        >> delete_test_dataset
     )
 
-create_test_dataset >> load_csv >> delete_test_dataset
+    from tests.system.utils.watcher import watcher
+
+    # This test needs watcher in order to properly mark success/failure
+    # when "tearDown" task with trigger rule is part of the DAG
+    list(dag.tasks) >> watcher()
+
+from tests.system.utils import get_test_run  # noqa: E402
+
+# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
+test_run = get_test_run(dag)
