diff --git a/airflow/providers/google/cloud/example_dags/example_automl_video_intelligence_classification.py b/tests/system/providers/google/cloud/automl/example_automl_video_intelligence_classification.py similarity index 79% rename from airflow/providers/google/cloud/example_dags/example_automl_video_intelligence_classification.py rename to tests/system/providers/google/cloud/automl/example_automl_video_intelligence_classification.py index f98be55bb0e39d..79314453060cb0 100644 --- a/airflow/providers/google/cloud/example_dags/example_automl_video_intelligence_classification.py +++ b/tests/system/providers/google/cloud/automl/example_automl_video_intelligence_classification.py @@ -32,6 +32,9 @@ AutoMLTrainModelOperator, ) +ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID") +DAG_ID = "example_automl_video" + GCP_PROJECT_ID = os.environ.get("GCP_PROJECT_ID", "your-project-id") GCP_AUTOML_LOCATION = os.environ.get("GCP_AUTOML_LOCATION", "us-central1") GCP_AUTOML_VIDEO_BUCKET = os.environ.get( @@ -61,12 +64,13 @@ # Example DAG for AutoML Video Intelligence Classification with models.DAG( - "example_automl_video", + DAG_ID, + schedule="@once", # Override to match your needs start_date=datetime(2021, 1, 1), catchup=False, user_defined_macros={"extract_object_id": extract_object_id}, - tags=['example'], -) as example_dag: + tags=['example', 'automl'], +) as dag: create_dataset_task = AutoMLCreateDatasetOperator( task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION ) @@ -100,11 +104,24 @@ project_id=GCP_PROJECT_ID, ) - import_dataset_task >> create_model - delete_model_task >> delete_datasets_task + ( + # TEST SETUP + create_dataset_task + >> import_dataset_task + # TEST BODY + >> create_model + # TEST TEARDOWN + >> delete_model_task + >> delete_datasets_task + ) + + 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 - # Task dependencies created via `XComArgs`: - # create_dataset_task >> import_dataset_task - # create_dataset_task >> create_model - # create_model >> delete_model_task - # create_dataset_task >> delete_datasets_task +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag)