diff --git a/airflow/providers/google/cloud/example_dags/example_automl_nl_text_classification.py b/tests/system/providers/google/cloud/automl/example_automl_nl_text_classification.py similarity index 87% rename from airflow/providers/google/cloud/example_dags/example_automl_nl_text_classification.py rename to tests/system/providers/google/cloud/automl/example_automl_nl_text_classification.py index afe28a9ff82ec..39129162097e0 100644 --- a/airflow/providers/google/cloud/example_dags/example_automl_nl_text_classification.py +++ b/tests/system/providers/google/cloud/automl/example_automl_nl_text_classification.py @@ -39,13 +39,9 @@ GCP_AUTOML_LOCATION = os.environ.get("GCP_AUTOML_LOCATION", "us-central1") GCP_AUTOML_TEXT_CLS_BUCKET = os.environ.get("GCP_AUTOML_TEXT_CLS_BUCKET", "gs://INVALID BUCKET NAME") -# Example values -DATASET_ID = "" - # Example model MODEL = { "display_name": "auto_model_1", - "dataset_id": DATASET_ID, "text_classification_model_metadata": {}, } @@ -55,6 +51,7 @@ "text_classification_dataset_metadata": {"classification_type": "MULTICLASS"}, } + IMPORT_INPUT_CONFIG = {"gcs_source": {"input_uris": [GCP_AUTOML_TEXT_CLS_BUCKET]}} extract_object_id = CloudAutoMLHook.extract_object_id @@ -65,12 +62,13 @@ start_date=datetime(2021, 1, 1), catchup=False, tags=["example"], -) as example_dag: +) as dag: create_dataset_task = AutoMLCreateDatasetOperator( task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION ) dataset_id = cast(str, XComArg(create_dataset_task, key="dataset_id")) + MODEL["dataset_id"] = dataset_id import_dataset_task = AutoMLImportDataOperator( task_id="import_dataset_task", @@ -78,11 +76,9 @@ location=GCP_AUTOML_LOCATION, input_config=IMPORT_INPUT_CONFIG, ) - MODEL["dataset_id"] = dataset_id create_model = AutoMLTrainModelOperator(task_id="create_model", model=MODEL, location=GCP_AUTOML_LOCATION) - model_id = cast(str, XComArg(create_model, key="model_id")) delete_model_task = AutoMLDeleteModelOperator( @@ -99,10 +95,23 @@ project_id=GCP_PROJECT_ID, ) + # TEST BODY import_dataset_task >> create_model + # TEST TEARDOWN delete_model_task >> delete_datasets_task # Task dependencies created via `XComArgs`: # create_dataset_task >> import_dataset_task # create_dataset_task >> create_model # create_dataset_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 + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag) diff --git a/airflow/providers/google/cloud/example_dags/example_automl_nl_text_sentiment.py b/tests/system/providers/google/cloud/automl/example_automl_nl_text_sentiment.py similarity index 88% rename from airflow/providers/google/cloud/example_dags/example_automl_nl_text_sentiment.py rename to tests/system/providers/google/cloud/automl/example_automl_nl_text_sentiment.py index a823b8af89538..6bbe0d2bfbea8 100644 --- a/airflow/providers/google/cloud/example_dags/example_automl_nl_text_sentiment.py +++ b/tests/system/providers/google/cloud/automl/example_automl_nl_text_sentiment.py @@ -39,13 +39,9 @@ GCP_AUTOML_LOCATION = os.environ.get("GCP_AUTOML_LOCATION", "us-central1") GCP_AUTOML_SENTIMENT_BUCKET = os.environ.get("GCP_AUTOML_SENTIMENT_BUCKET", "gs://INVALID BUCKET NAME") -# Example values -DATASET_ID = "" - # Example model MODEL = { "display_name": "auto_model_1", - "dataset_id": DATASET_ID, "text_sentiment_model_metadata": {}, } @@ -66,12 +62,13 @@ catchup=False, user_defined_macros={"extract_object_id": extract_object_id}, tags=["example"], -) as example_dag: +) as dag: create_dataset_task = AutoMLCreateDatasetOperator( task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION ) dataset_id = cast(str, XComArg(create_dataset_task, key="dataset_id")) + MODEL["dataset_id"] = dataset_id import_dataset_task = AutoMLImportDataOperator( task_id="import_dataset_task", @@ -100,7 +97,9 @@ project_id=GCP_PROJECT_ID, ) + # TEST BODY import_dataset_task >> create_model + # TEST TEARDOWN delete_model_task >> delete_datasets_task # Task dependencies created via `XComArgs`: @@ -108,3 +107,14 @@ # create_dataset_task >> create_model # create_model >> delete_model_task # create_dataset_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 + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag) diff --git a/airflow/providers/google/cloud/example_dags/example_automl_translation.py b/tests/system/providers/google/cloud/automl/example_automl_translation.py similarity index 88% rename from airflow/providers/google/cloud/example_dags/example_automl_translation.py rename to tests/system/providers/google/cloud/automl/example_automl_translation.py index 2bef20caabee3..87bf7166dcd1f 100644 --- a/airflow/providers/google/cloud/example_dags/example_automl_translation.py +++ b/tests/system/providers/google/cloud/automl/example_automl_translation.py @@ -41,13 +41,9 @@ "GCP_AUTOML_TRANSLATION_BUCKET", "gs://INVALID BUCKET NAME/file" ) -# Example values -DATASET_ID = "TRL123456789" - # Example model MODEL = { "display_name": "auto_model_1", - "dataset_id": DATASET_ID, "translation_model_metadata": {}, } @@ -60,6 +56,7 @@ }, } + IMPORT_INPUT_CONFIG = {"gcs_source": {"input_uris": [GCP_AUTOML_TRANSLATION_BUCKET]}} extract_object_id = CloudAutoMLHook.extract_object_id @@ -69,10 +66,11 @@ with models.DAG( "example_automl_translation", start_date=datetime(2021, 1, 1), + schedule="@once", catchup=False, user_defined_macros={"extract_object_id": extract_object_id}, tags=["example"], -) as example_dag: +) as dag: create_dataset_task = AutoMLCreateDatasetOperator( task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION ) @@ -106,7 +104,9 @@ project_id=GCP_PROJECT_ID, ) + # TEST BODY import_dataset_task >> create_model + # TEST TEARDOWN delete_model_task >> delete_datasets_task # Task dependencies created via `XComArgs`: @@ -114,3 +114,15 @@ # create_dataset_task >> create_model # create_model >> delete_model_task # create_dataset_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 + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag) 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 88% 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 2b55c42a8ab96..63447a123059d 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 @@ -41,13 +41,9 @@ "GCP_AUTOML_VIDEO_BUCKET", "gs://INVALID BUCKET NAME/hmdb_split1.csv" ) -# Example values -DATASET_ID = "VCN123455678" - # Example model MODEL = { "display_name": "auto_model_1", - "dataset_id": DATASET_ID, "video_classification_model_metadata": {}, } @@ -69,12 +65,13 @@ catchup=False, user_defined_macros={"extract_object_id": extract_object_id}, tags=["example"], -) as example_dag: +) as dag: create_dataset_task = AutoMLCreateDatasetOperator( task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION ) dataset_id = cast(str, XComArg(create_dataset_task, key="dataset_id")) + MODEL["dataset_id"] = dataset_id import_dataset_task = AutoMLImportDataOperator( task_id="import_dataset_task", @@ -103,7 +100,9 @@ project_id=GCP_PROJECT_ID, ) + # TEST BODY import_dataset_task >> create_model + # TEST TEARDOWN delete_model_task >> delete_datasets_task # Task dependencies created via `XComArgs`: @@ -111,3 +110,14 @@ # create_dataset_task >> create_model # create_model >> delete_model_task # create_dataset_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 + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag) diff --git a/airflow/providers/google/cloud/example_dags/example_automl_video_intelligence_tracking.py b/tests/system/providers/google/cloud/automl/example_automl_video_intelligence_tracking.py similarity index 88% rename from airflow/providers/google/cloud/example_dags/example_automl_video_intelligence_tracking.py rename to tests/system/providers/google/cloud/automl/example_automl_video_intelligence_tracking.py index daed5748b2399..5ac8fa7457fa6 100644 --- a/airflow/providers/google/cloud/example_dags/example_automl_video_intelligence_tracking.py +++ b/tests/system/providers/google/cloud/automl/example_automl_video_intelligence_tracking.py @@ -42,13 +42,10 @@ "gs://INVALID BUCKET NAME/youtube_8m_videos_animal_tiny.csv", ) -# Example values -DATASET_ID = "VOT123456789" # Example model MODEL = { "display_name": "auto_model_1", - "dataset_id": DATASET_ID, "video_object_tracking_model_metadata": {}, } @@ -70,12 +67,13 @@ catchup=False, user_defined_macros={"extract_object_id": extract_object_id}, tags=["example"], -) as example_dag: +) as dag: create_dataset_task = AutoMLCreateDatasetOperator( task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION ) dataset_id = cast(str, XComArg(create_dataset_task, key="dataset_id")) + MODEL["dataset_id"] = dataset_id import_dataset_task = AutoMLImportDataOperator( task_id="import_dataset_task", @@ -104,7 +102,9 @@ project_id=GCP_PROJECT_ID, ) + # TEST BODY import_dataset_task >> create_model + # TEST TEARDOWN delete_model_task >> delete_datasets_task # Task dependencies created via `XComArgs`: @@ -112,3 +112,14 @@ # create_dataset_task >> create_model # create_model >> delete_model_task # create_dataset_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 + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag) diff --git a/airflow/providers/google/cloud/example_dags/example_automl_vision_object_detection.py b/tests/system/providers/google/cloud/automl/example_automl_vision_object_detection.py similarity index 88% rename from airflow/providers/google/cloud/example_dags/example_automl_vision_object_detection.py rename to tests/system/providers/google/cloud/automl/example_automl_vision_object_detection.py index 8b9ae271d1126..933c2760cbbbd 100644 --- a/airflow/providers/google/cloud/example_dags/example_automl_vision_object_detection.py +++ b/tests/system/providers/google/cloud/automl/example_automl_vision_object_detection.py @@ -41,13 +41,9 @@ "GCP_AUTOML_DETECTION_BUCKET", "gs://INVALID BUCKET NAME/img/openimage/csv/salads_ml_use.csv" ) -# Example values -DATASET_ID = "" - # Example model MODEL = { "display_name": "auto_model", - "dataset_id": DATASET_ID, "image_object_detection_model_metadata": {}, } @@ -69,12 +65,13 @@ catchup=False, user_defined_macros={"extract_object_id": extract_object_id}, tags=["example"], -) as example_dag: +) as dag: create_dataset_task = AutoMLCreateDatasetOperator( task_id="create_dataset_task", dataset=DATASET, location=GCP_AUTOML_LOCATION ) dataset_id = cast(str, XComArg(create_dataset_task, key="dataset_id")) + MODEL["dataset_id"] = dataset_id import_dataset_task = AutoMLImportDataOperator( task_id="import_dataset_task", @@ -103,7 +100,9 @@ project_id=GCP_PROJECT_ID, ) + # TEST BODY import_dataset_task >> create_model + # TEST TEARDOWN delete_model_task >> delete_datasets_task # Task dependencies created via `XComArgs`: @@ -111,3 +110,14 @@ # create_dataset_task >> create_model # create_model >> delete_model_task # create_dataset_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 + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag)