AIP-47 - Migrate livy DAGs to new design #22439 (#24208)

apache · Jun 5, 2022 · 5503603 · 5503603
1 parent ac8a790
commit 5503603
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 19 deletions.
diff --git a/airflow/providers/apache/livy/example_dags/__init__.py b/airflow/providers/apache/livy/example_dags/__init__.py
diff --git a/docs/apache-airflow-providers-apache-livy/index.rst b/docs/apache-airflow-providers-apache-livy/index.rst
@@ -37,7 +37,7 @@ Content
     :maxdepth: 1
     :caption: Resources
 
-    Example DAGs <https://github.com/apache/airflow/tree/main/airflow/providers/apache/livy/example_dags>
+    Example DAGs <https://github.com/apache/airflow/tree/main/tests/system/providers/apache/livy>
     PyPI Repository <https://pypi.org/project/apache-airflow-providers-apache-livy/>
     Installing from sources <installing-providers-from-sources>
 

diff --git a/docs/apache-airflow-providers-apache-livy/operators.rst b/docs/apache-airflow-providers-apache-livy/operators.rst
@@ -29,7 +29,7 @@ LivyOperator
 
 This operator wraps the Apache Livy batch REST API, allowing to submit a Spark application to the underlying cluster.
 
-.. exampleinclude:: /../../airflow/providers/apache/livy/example_dags/example_livy.py
+.. exampleinclude:: /../../tests/system/providers/apache/livy/example_livy.py
     :language: python
     :start-after: [START create_livy]
     :end-before: [END create_livy]

diff --git a/.../apache/livy/example_dags/example_livy.py → ...tem/providers/apache/livy/example_livy.py b/.../apache/livy/example_dags/example_livy.py → ...tem/providers/apache/livy/example_livy.py
@@ -20,13 +20,18 @@
 The tasks below trigger the computation of pi on the Spark instance
 using the Java and Python executables provided in the example library.
 """
+
+import os
 from datetime import datetime
 
 from airflow import DAG
 from airflow.providers.apache.livy.operators.livy import LivyOperator
 
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "example_livy_operator"
+
 with DAG(
-    dag_id='example_livy_operator',
+    dag_id=DAG_ID,
     default_args={'args': [10]},
     schedule_interval='@daily',
     start_date=datetime(2021, 1, 1),
@@ -48,3 +53,14 @@
 
     livy_java_task >> livy_python_task
     # [END create_livy]
+
+    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)