Migrate Microsoft example DAGs to new design #22452 - azure

docs/apache-airflow-providers-microsoft-azure/index.rst

@@ -41,7 +41,7 @@ Content
     :maxdepth: 1
     :caption: Resources
 
-    Example DAGs <https://github.com/apache/airflow/tree/main/airflow/providers/microsoft/azure/example_dags>
+    Example DAGs <https://github.com/apache/airflow/tree/main/tests/system/providers/microsoft/azure>
     PyPI Repository <https://pypi.org/project/apache-airflow-providers-microsoft-azure/>
     Installing from sources <installing-providers-from-sources>
 

docs/apache-airflow-providers-microsoft-azure/operators/adf_run_pipeline.rst

@@ -31,15 +31,15 @@ This functionality can be disabled for an asynchronous wait -- typically with th
 
 Below is an example of using this operator to execute an Azure Data Factory pipeline.
 
-  .. exampleinclude:: /../../airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py
+  .. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_adf_run_pipeline.py
       :language: python
       :dedent: 0
       :start-after: [START howto_operator_adf_run_pipeline]
       :end-before: [END howto_operator_adf_run_pipeline]
 
 Here is a different example of using this operator to execute a pipeline but coupled with the :class:`~airflow.providers.microsoft.azure.sensors.data_factory.AzureDataFactoryPipelineRunSensor` to perform an asynchronous wait.
 
-    .. exampleinclude:: /../../airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py
+    .. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_adf_run_pipeline.py
         :language: python
         :dedent: 0
         :start-after: [START howto_operator_adf_run_pipeline_async]

docs/apache-airflow-providers-microsoft-azure/operators/adls.rst

@@ -35,7 +35,7 @@ file(s) from Azure DataLake Storage
 
 Below is an example of using this operator to delete a file from ADL.
 
-.. exampleinclude:: /../../airflow/providers/microsoft/azure/example_dags/example_adls_delete.py
+.. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_adls_delete.py
     :language: python
     :dedent: 0
     :start-after: [START howto_operator_adls_delete]

docs/apache-airflow-providers-microsoft-azure/operators/azure_blob_to_gcs.rst

@@ -54,7 +54,7 @@ To get information about jobs within a Azure Blob Storage use:
 
 Example usage:
 
-.. exampleinclude:: /../../airflow/providers/microsoft/azure/example_dags/example_azure_blob_to_gcs.py
+.. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_azure_blob_to_gcs.py
     :language: python
     :start-after: [START how_to_azure_blob_to_gcs]
     :end-before: [END how_to_azure_blob_to_gcs]

docs/apache-airflow-providers-microsoft-azure/operators/local_to_adls.rst

@@ -38,7 +38,7 @@ upload data from local filesystem to ADL.
 
 Below is an example of using this operator to upload a file to ADL.
 
-.. exampleinclude:: /../../airflow/providers/microsoft/azure/example_dags/example_local_to_adls.py
+.. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_local_to_adls.py
     :language: python
     :dedent: 0
     :start-after: [START howto_operator_local_to_adls]

docs/apache-airflow-providers-microsoft-azure/operators/sftp_to_wasb.rst

@@ -54,7 +54,7 @@ To get information about jobs within a Azure Blob Storage use:
 :class:`~airflow.providers.microsoft.azure.transfers.sftp_to_wasb.SFTPToWasbOperator`
 Example usage:
 
-.. exampleinclude:: /../../airflow/providers/microsoft/azure/example_dags/example_sftp_to_wasb.py
+.. exampleinclude:: /../../tests/system/providers/microsoft/azure/example_sftp_to_wasb.py
     :language: python
     :dedent: 4
     :start-after: [START how_to_sftp_to_wasb]

tests/providers/microsoft/azure/operators/test_adls_delete_system.py

@@ -19,7 +19,7 @@
 
 import pytest
 
-from airflow.providers.microsoft.azure.example_dags.example_local_to_adls import LOCAL_FILE_PATH
+from tests.system.providers.microsoft.azure.example_local_to_adls import LOCAL_FILE_PATH
 from tests.test_utils.azure_system_helpers import (
     AZURE_DAG_FOLDER,
     AzureSystemTest,

tests/providers/microsoft/azure/transfers/test_local_to_adls_system.py

@@ -19,7 +19,7 @@
 
 import pytest
 
-from airflow.providers.microsoft.azure.example_dags.example_local_to_adls import LOCAL_FILE_PATH
+from tests.system.providers.microsoft.azure.example_local_to_adls import LOCAL_FILE_PATH
 from tests.test_utils.azure_system_helpers import (
     AZURE_DAG_FOLDER,
     AzureSystemTest,

tests/providers/microsoft/azure/transfers/test_local_to_wasb_system.py

@@ -19,9 +19,7 @@
 
 import pytest
 
-from airflow.providers.microsoft.azure.example_dags.example_local_to_wasb import (  # type: ignore
-    PATH_TO_UPLOAD_FILE,
-)
+from tests.system.providers.microsoft.azure.example_local_to_wasb import PATH_TO_UPLOAD_FILE  # type: ignore
 from tests.test_utils.azure_system_helpers import (
     AZURE_DAG_FOLDER,
     AzureSystemTest,

tests/providers/microsoft/azure/transfers/test_sftp_to_wasb_system.py

@@ -20,7 +20,7 @@
 
 import pytest
 
-from airflow.providers.microsoft.azure.example_dags.example_sftp_to_wasb import (
+from tests.system.providers.microsoft.azure.example_sftp_to_wasb import (
     FILE_COMPLETE_PATH,
     LOCAL_FILE_PATH,
     SAMPLE_FILENAME,

tests/system/providers/__init__.py

@@ -0,0 +1,17 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

tests/system/providers/microsoft/__init__.py

@@ -0,0 +1,17 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

tests/system/providers/microsoft/azure/__init__.py

@@ -0,0 +1,17 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

airflow/providers/microsoft/azure/example_dags/example_adf_run_pipeline.py → tests/system/providers/microsoft/azure/example_adf_run_pipeline.py

@@ -14,21 +14,25 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-
+import os
 from datetime import datetime, timedelta
 
-from airflow.models import DAG, BaseOperator
+from airflow.models import DAG
 
 try:
     from airflow.operators.empty import EmptyOperator
 except ModuleNotFoundError:
     from airflow.operators.dummy import DummyOperator as EmptyOperator  # type: ignore
+
 from airflow.providers.microsoft.azure.operators.data_factory import AzureDataFactoryRunPipelineOperator
 from airflow.providers.microsoft.azure.sensors.data_factory import AzureDataFactoryPipelineRunStatusSensor
 from airflow.utils.edgemodifier import Label
 
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "example_adf_run_pipeline"
+
 with DAG(
-    dag_id="example_adf_run_pipeline",
+    dag_id=DAG_ID,
     start_date=datetime(2021, 8, 13),
     schedule_interval="@daily",
     catchup=False,
@@ -45,21 +49,21 @@
     end = EmptyOperator(task_id="end")
 
     # [START howto_operator_adf_run_pipeline]
-    run_pipeline1: BaseOperator = AzureDataFactoryRunPipelineOperator(
+    run_pipeline1 = AzureDataFactoryRunPipelineOperator(
         task_id="run_pipeline1",
         pipeline_name="pipeline1",
         parameters={"myParam": "value"},
     )
     # [END howto_operator_adf_run_pipeline]
 
     # [START howto_operator_adf_run_pipeline_async]
-    run_pipeline2: BaseOperator = AzureDataFactoryRunPipelineOperator(
+    run_pipeline2 = AzureDataFactoryRunPipelineOperator(
         task_id="run_pipeline2",
         pipeline_name="pipeline2",
         wait_for_termination=False,
     )
 
-    pipeline_run_sensor: BaseOperator = AzureDataFactoryPipelineRunStatusSensor(
+    pipeline_run_sensor = AzureDataFactoryPipelineRunStatusSensor(
         task_id="pipeline_run_sensor",
         run_id=run_pipeline2.output["run_id"],
     )
@@ -71,3 +75,14 @@
 
     # Task dependency created via `XComArgs`:
     #   run_pipeline2 >> pipeline_run_sensor
+
+    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)

airflow/providers/microsoft/azure/example_dags/example_adls_delete.py → tests/system/providers/microsoft/azure/example_adls_delete.py

@@ -25,9 +25,11 @@
 LOCAL_FILE_PATH = os.environ.get("LOCAL_FILE_PATH", 'localfile.txt')
 REMOTE_FILE_PATH = os.environ.get("REMOTE_LOCAL_PATH", 'remote.txt')
 
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "example_adls_delete"
 
 with models.DAG(
-    "example_adls_delete",
+    DAG_ID,
     start_date=datetime(2021, 1, 1),
     schedule_interval=None,
     tags=['example'],
@@ -43,3 +45,14 @@
     # [END howto_operator_adls_delete]
 
     upload_file >> remove_file
+
+    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)

airflow/providers/microsoft/azure/example_dags/example_azure_blob_to_gcs.py → tests/system/providers/microsoft/azure/example_azure_blob_to_gcs.py

@@ -31,10 +31,12 @@
 GCP_BUCKET_FILE_PATH = os.environ.get("GCP_BUCKET_FILE_PATH", "file.txt")
 GCP_BUCKET_NAME = os.environ.get("GCP_BUCKET_NAME", "INVALID BUCKET NAME")
 GCP_OBJECT_NAME = os.environ.get("GCP_OBJECT_NAME", "file.txt")
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "example_azure_blob_to_gcs"
 
 # [START how_to_azure_blob_to_gcs]
 with DAG(
-    "example_azure_blob_to_gcs",
+    DAG_ID,
     schedule_interval=None,
     start_date=datetime(2021, 1, 1),  # Override to match your needs
     default_args={"container_name": AZURE_CONTAINER_NAME, "blob_name": BLOB_NAME},
@@ -57,3 +59,14 @@
     # [END how_to_azure_blob_to_gcs]
 
     wait_for_blob >> transfer_files_to_gcs
+
+    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)

airflow/providers/microsoft/azure/example_dags/example_azure_container_instances.py → tests/system/providers/microsoft/azure/example_azure_container_instances.py

@@ -18,13 +18,17 @@
 """
 This is an example dag for using the AzureContainerInstancesOperator.
 """
+import os
 from datetime import datetime, timedelta
 
 from airflow import DAG
 from airflow.providers.microsoft.azure.operators.container_instances import AzureContainerInstancesOperator
 
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "aci_example"
+
 with DAG(
-    dag_id='aci_example',
+    dag_id=DAG_ID,
     default_args={'retries': 1},
     schedule_interval=timedelta(days=1),
     start_date=datetime(2018, 11, 1),
@@ -45,3 +49,9 @@
         cpu=1.0,
         task_id='start_container',
     )
+
+
+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)

airflow/providers/microsoft/azure/example_dags/example_azure_cosmosdb.py → tests/system/providers/microsoft/azure/example_azure_cosmosdb.py

@@ -28,15 +28,18 @@
 *Note: Make sure that connection `azure_cosmos_default` is properly set before running
 this example.*
 """
-
+import os
 from datetime import datetime
 
 from airflow import DAG
 from airflow.providers.microsoft.azure.operators.cosmos import AzureCosmosInsertDocumentOperator
 from airflow.providers.microsoft.azure.sensors.cosmos import AzureCosmosDocumentSensor
 
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "example_azure_cosmosdb_sensor"
+
 with DAG(
-    dag_id='example_azure_cosmosdb_sensor',
+    dag_id=DAG_ID,
     default_args={'database_name': 'airflow_example_db'},
     start_date=datetime(2021, 1, 1),
     catchup=False,
@@ -57,3 +60,14 @@
     )
 
     t1 >> t2
+
+    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)

airflow/providers/microsoft/azure/example_dags/example_fileshare.py → tests/system/providers/microsoft/azure/example_fileshare.py

@@ -14,7 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-
+import os
 from datetime import datetime
 
 from airflow.decorators import task
@@ -23,6 +23,8 @@
 
 NAME = 'myfileshare'
 DIRECTORY = "mydirectory"
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "example_fileshare"
 
 
 @task
@@ -44,9 +46,20 @@ def delete_fileshare():
 
 
 with DAG(
-    "example_fileshare",
+    DAG_ID,
     schedule_interval="@once",
     start_date=datetime(2021, 1, 1),
     catchup=False,
 ) as dag:
     create_fileshare() >> delete_fileshare()
+
+    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)