From 0e11fa32c522773ae1a1e7b6a72ef61e7876731d Mon Sep 17 00:00:00 2001
From: Chethan UK <chethan.umesha@tessian.com>
Date: Fri, 3 Jun 2022 01:16:14 +0100
Subject: [PATCH] Migrate Plexus example DAGs to new design #22457

---
 .../providers/plexus/example_dags/__init__.py | 16 --------
 .../apache-airflow-providers-plexus/index.rst |  2 +-
 .../providers/plexus}/example_plexus.py       | 40 +++++++++++--------
 3 files changed, 25 insertions(+), 33 deletions(-)
 delete mode 100644 airflow/providers/plexus/example_dags/__init__.py
 rename {airflow/providers/plexus/example_dags => tests/system/providers/plexus}/example_plexus.py (65%)

diff --git a/airflow/providers/plexus/example_dags/__init__.py b/airflow/providers/plexus/example_dags/__init__.py
deleted file mode 100644
index 13a83393a9124b..00000000000000
--- a/airflow/providers/plexus/example_dags/__init__.py
+++ /dev/null
@@ -1,16 +0,0 @@
-# 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.
diff --git a/docs/apache-airflow-providers-plexus/index.rst b/docs/apache-airflow-providers-plexus/index.rst
index 8926d9e67e2666..df86306ba7e445 100644
--- a/docs/apache-airflow-providers-plexus/index.rst
+++ b/docs/apache-airflow-providers-plexus/index.rst
@@ -32,7 +32,7 @@ Content
     :maxdepth: 1
     :caption: Resources
 
-    Example DAGs <https://github.com/apache/airflow/tree/main/airflow/providers/plexus/example_dags>
+    Example DAGs <https://github.com/apache/airflow/tree/main/tests/system/providers/plexus>
 
 .. toctree::
     :maxdepth: 1
diff --git a/airflow/providers/plexus/example_dags/example_plexus.py b/tests/system/providers/plexus/example_plexus.py
similarity index 65%
rename from airflow/providers/plexus/example_dags/example_plexus.py
rename to tests/system/providers/plexus/example_plexus.py
index 68ddcb7d030d2c..17c7f499c0aa00 100644
--- a/airflow/providers/plexus/example_dags/example_plexus.py
+++ b/tests/system/providers/plexus/example_plexus.py
@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import os
 from datetime import datetime
 
 from airflow import DAG
@@ -22,26 +23,33 @@
 
 HOME = '/home/acc'
 T3_PRERUN_SCRIPT = 'cp {home}/imdb/run_scripts/mlflow.sh {home}/ && chmod +x mlflow.sh'.format(home=HOME)
+ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
+DAG_ID = "test"
 
-
-dag = DAG(
-    'test',
+with DAG(
+    DAG_ID,
     default_args={'owner': 'core scientific', 'retries': 1},
     description='testing plexus operator',
     start_date=datetime(2021, 1, 1),
     schedule_interval='@once',
     catchup=False,
-)
+) as dag:
+    # [START plexus_job_op]
+    t1 = PlexusJobOperator(
+        task_id='test',
+        job_params={
+            'name': 'test',
+            'app': 'MLFlow Pipeline 01',
+            'queue': 'DGX-2 (gpu:Tesla V100-SXM3-32GB)',
+            'num_nodes': 1,
+            'num_cores': 1,
+            'prerun_script': T3_PRERUN_SCRIPT,
+        },
+    )
+    # [END plexus_job_op]
+
+
+from tests.system.utils import get_test_run  # noqa: E402
 
-t1 = PlexusJobOperator(
-    task_id='test',
-    job_params={
-        'name': 'test',
-        'app': 'MLFlow Pipeline 01',
-        'queue': 'DGX-2 (gpu:Tesla V100-SXM3-32GB)',
-        'num_nodes': 1,
-        'num_cores': 1,
-        'prerun_script': T3_PRERUN_SCRIPT,
-    },
-    dag=dag,
-)
+# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)
+test_run = get_test_run(dag)