Repository: PrefectHQ/prefect

Issue:
refactor airflow tests
We should always be testing (and probably even expanding our tests) for the airflow compatibility layer.  Right now, #246 introduces a new flag for `pytest` which optionally runs the airflow tests (which are very slow, around ~2min).

We should:
- refactor `AirFlow` to allow for using an already live airflow database
- refactor the tests to rely on a single fixed airflow db instead of recreating one with each test
- create a separate airflow CI job to run these tests so they don't slow down PRs too much


Discussion:

Resolving patch:
diff --git a/src/prefect/utilities/airflow_utils.py b/src/prefect/utilities/airflow_utils.py
--- a/src/prefect/utilities/airflow_utils.py
+++ b/src/prefect/utilities/airflow_utils.py
@@ -30,6 +30,13 @@
 __all__ = ["AirFlow"]
 
 
+def custom_query(db: str, query: str) -> List:
+    with closing(sqlite3.connect(db)) as connection:
+        with closing(connection.cursor()) as cursor:
+            cursor.execute(query)
+            return cursor.fetchall()
+
+
 trigger_mapping = {
     "all_success": prefect.triggers.all_successful,
     "all_failed": prefect.triggers.all_failed,
@@ -40,13 +47,6 @@
 }
 
 
-def custom_query(db: str, query: str) -> List:
-    with closing(sqlite3.connect(db)) as connection:
-        with closing(connection.cursor()) as cursor:
-            cursor.execute(query)
-            return cursor.fetchall()
-
-
 class AirTask(prefect.tasks.shell.ShellTask):
     def __init__(self, task: "airflow.models.BaseOperator", **kwargs: Any) -> None:
         name = task.task_id
@@ -58,10 +58,10 @@ def __init__(self, task: "airflow.models.BaseOperator", **kwargs: Any) -> None:
         super().__init__(
             name=name,
             command=cmd,
-            trigger=trigger,
             skip_on_upstream_skip=False,
             max_retries=task.retries,
             tags=task.pool,
+            trigger=trigger,
             **kwargs
         )
 
@@ -158,19 +158,31 @@ class AirFlow(prefect.core.flow.Flow):
         ```
     """
 
-    def __init__(self, dag_id: str, *args: Any, **kwargs: Any) -> None:
-        self.dag = airflow.models.DagBag().dags[dag_id]
+    def __init__(
+        self,
+        dag_id: str,
+        *args: Any,
+        db_file: str = None,
+        dag_folder: str = None,
+        **kwargs: Any
+    ) -> None:
+        self.dag = airflow.models.DagBag(dag_folder=dag_folder).dags[dag_id]
         super().__init__(*args, **kwargs)
         self._populate_tasks()
-        self.env = self._init_db()
+        self.env = self._init_db(db_file=db_file, dag_folder=dag_folder)
 
-    def _init_db(self) -> dict:
-        self.td = tempfile.TemporaryDirectory(prefix="prefect-airflow")
+    def _init_db(self, db_file: str = None, dag_folder: str = None) -> dict:
         env = os.environ.copy()
-        env["AIRFLOW__CORE__SQL_ALCHEMY_CONN"] = (
-            "sqlite:///" + self.td.name + "/prefect-airflow.db"
-        )
-        status = subprocess.check_output(["bash", "-c", "airflow initdb"], env=env)
+        if dag_folder is not None:
+            env["AIRFLOW__CORE__DAGS_FOLDER"] = dag_folder
+        if db_file is not None:
+            env["AIRFLOW__CORE__SQL_ALCHEMY_CONN"] = "sqlite:///" + db_file
+        else:
+            self.td = tempfile.TemporaryDirectory(prefix="prefect-airflow")
+            env["AIRFLOW__CORE__SQL_ALCHEMY_CONN"] = (
+                "sqlite:///" + self.td.name + "/prefect-airflow.db"
+            )
+            status = subprocess.check_output(["bash", "-c", "airflow initdb"], env=env)
         return env
 
     def run(  # type: ignore


Test patch:
diff --git a/tests/utilities/airflow/dags/trigger_rules.py b/tests/utilities/airflow/dags/trigger_rules.py
new file mode 100644
--- /dev/null
+++ b/tests/utilities/airflow/dags/trigger_rules.py
@@ -0,0 +1,17 @@
+import datetime as dt
+
+import airflow
+from airflow.models import DAG
+from airflow.operators.dummy_operator import DummyOperator
+from airflow.utils.trigger_rule import TriggerRule
+
+dag = DAG(
+    dag_id="trigger_rules",
+    schedule_interval=dt.timedelta(hours=4),
+    start_date=airflow.utils.dates.days_ago(2),
+)
+
+task1 = DummyOperator(task_id="task1", dag=dag)
+task2 = DummyOperator(task_id="task2", dag=dag)
+task3 = DummyOperator(task_id="task3", dag=dag, trigger_rule=TriggerRule.ONE_FAILED)
+task3.set_upstream([task1, task2])
diff --git a/tests/utilities/airflow/test_airflow.py b/tests/utilities/airflow/test_airflow.py
new file mode 100644
--- /dev/null
+++ b/tests/utilities/airflow/test_airflow.py
@@ -0,0 +1,139 @@
+import os
+import pytest
+import subprocess
+import tempfile
+
+airflow = pytest.importorskip("airflow")
+
+from prefect import task, Task, triggers
+from prefect.engine.state import Success, Skipped, Failed, TriggerFailed
+from prefect.utilities.airflow_utils import AirFlow
+
+
+@pytest.fixture(scope="module")
+def airflow_settings():
+    with tempfile.NamedTemporaryFile(
+        prefix="prefect-airflow", suffix="prefect-airflow-test.db"
+    ) as tmp:
+        env = os.environ.copy()
+        env["AIRFLOW__CORE__SQL_ALCHEMY_CONN"] = "sqlite:///" + tmp.name
+        dag_folder = os.path.join(os.path.dirname(__file__), "dags")
+        env["AIRFLOW__CORE__DAGS_FOLDER"] = dag_folder
+        status = subprocess.check_output(["bash", "-c", "airflow initdb"], env=env)
+        yield dict(db_file=tmp.name, dag_folder=dag_folder)
+
+
+def test_trigger_rules_dag(airflow_settings):
+    flow = AirFlow(dag_id="trigger_rules", **airflow_settings)
+    res = flow.run(execution_date="2018-09-20", return_tasks=flow.tasks)
+    assert res.is_failed()
+    for task, state in res.result.items():
+        if "3" in task.name:
+            assert isinstance(state, TriggerFailed)
+        else:
+            assert state.is_successful()
+
+
+@pytest.mark.airflow()
+def test_airflow_accepts_existing_sqlite_db():
+    flow = AirFlow(dag_id="example_bash_operator", db_file="test_doesnt_exist")
+    assert (
+        flow.env.get("AIRFLOW__CORE__SQL_ALCHEMY_CONN") == "sqlite:///test_doesnt_exist"
+    )
+
+
+@pytest.mark.airflow()
+def test_airflow_creates_sqlite_db_if_none_provided():
+    flow = AirFlow(dag_id="example_bash_operator")
+    sql_conn = flow.env.get("AIRFLOW__CORE__SQL_ALCHEMY_CONN")
+    assert sql_conn.startswith("sqlite:///")
+    assert sql_conn.endswith("prefect-airflow.db")
+
+
+@pytest.mark.airflow()
+def test_example_branch_operator(airflow_settings):
+    flow = AirFlow(dag_id="example_branch_operator", **airflow_settings)
+    res = flow.run(execution_date="2018-09-20", return_tasks=flow.tasks)
+    assert res.is_successful()
+
+    branch_tasks, follow_tasks = [], []
+    for task, state in res.result.items():
+        if task.name in ["run_this_first", "branching", "join"]:
+            assert state.is_successful()
+            assert state.result is None
+        elif "follow" in task.name:
+            follow_tasks.append((task, state))
+        elif "branch" in task.name:
+            branch_tasks.append((task, state))
+
+    run_branches = list(filter(lambda x: not isinstance(x[1], Skipped), branch_tasks))
+    assert len(run_branches) == 1
+    branch, state = run_branches.pop()
+    assert state.is_successful()
+    assert all([isinstance(s, Skipped) for t, s in branch_tasks if t != branch])
+    assert all(
+        [isinstance(s, Skipped) for t, s in follow_tasks if (branch.name not in t.name)]
+    )
+    assert all(
+        [isinstance(s, Success) for t, s in follow_tasks if (branch.name in t.name)]
+    )
+
+
+@pytest.mark.airflow()
+def test_example_xcom(airflow_settings):
+    flow = AirFlow(dag_id="example_xcom", **airflow_settings)
+    res = flow.run(execution_date="2018-09-20", return_tasks=flow.tasks)
+    assert res.is_successful()
+
+    for task, state in res.result.items():
+        if task.name == "puller":
+            assert state.is_successful()
+            assert state.result is None
+        elif task.name == "push":
+            assert state.is_successful()
+            assert state.result == [1, 2, 3]
+        if task.name == "push_by_returning":
+            assert state.is_successful()
+            assert state.result == {"a": "b"}
+
+
+@pytest.mark.airflow()
+def test_example_short_circuit_operator(airflow_settings):
+    flow = AirFlow(dag_id="example_short_circuit_operator", **airflow_settings)
+    res = flow.run(execution_date="2018-09-20", return_tasks=flow.tasks)
+    assert res.is_successful()
+
+    for task, state in res.result.items():
+        if "condition" in task.name:
+            assert isinstance(state, Success)
+        elif "false" in task.name:
+            assert isinstance(state, Skipped)
+        elif "true" in task.name:
+            assert isinstance(state, Success)
+
+
+@pytest.mark.airflow()
+def test_example_bash_operator(airflow_settings):
+    flow = AirFlow(dag_id="example_bash_operator", **airflow_settings)
+    res = flow.run(execution_date="2018-09-20", return_tasks=flow.tasks)
+    assert res.is_successful()
+
+    for task, state in res.result.items():
+        assert isinstance(state, Success)
+
+
+@pytest.mark.airflow()
+def test_extending_airflow_dag_with_prefect_task(airflow_settings):
+    flow = AirFlow(dag_id="example_bash_operator", **airflow_settings)
+    run_task0 = flow.get_tasks(name="runme_0")
+    with flow:
+        t1 = Task(trigger=triggers.all_failed)(upstream_tasks=[run_task0])
+
+    res = flow.run(execution_date="2018-09-21", return_tasks=flow.tasks)
+    assert res.is_failed()
+
+    for task, state in res.result.items():
+        if task == t1:
+            assert isinstance(state, TriggerFailed)
+        else:
+            assert isinstance(state, Success)
diff --git a/tests/utilities/test_airflow.py b/tests/utilities/test_airflow.py
deleted file mode 100644
--- a/tests/utilities/test_airflow.py
+++ /dev/null
@@ -1,78 +0,0 @@
-import pytest
-
-airflow = pytest.importorskip("airflow")
-
-from prefect.engine.state import Success, Skipped, Failed
-from prefect.utilities.airflow_utils import AirFlow
-
-
-@pytest.mark.airflow()
-def test_example_branch_operator():
-    flow = AirFlow(dag_id="example_branch_operator")
-    res = flow.run(execution_date="2018-09-20", return_tasks=flow.tasks)
-    assert res.is_successful()
-
-    branch_tasks, follow_tasks = [], []
-    for task, state in res.result.items():
-        if task.name in ["run_this_first", "branching", "join"]:
-            assert state.is_successful()
-            assert state.result is None
-        elif "follow" in task.name:
-            follow_tasks.append((task, state))
-        elif "branch" in task.name:
-            branch_tasks.append((task, state))
-
-    run_branches = list(filter(lambda x: not isinstance(x[1], Skipped), branch_tasks))
-    assert len(run_branches) == 1
-    branch, state = run_branches.pop()
-    assert state.is_successful()
-    assert all([isinstance(s, Skipped) for t, s in branch_tasks if t != branch])
-    assert all(
-        [isinstance(s, Skipped) for t, s in follow_tasks if (branch.name not in t.name)]
-    )
-    assert all(
-        [isinstance(s, Success) for t, s in follow_tasks if (branch.name in t.name)]
-    )
-
-
-@pytest.mark.airflow()
-def test_example_xcom():
-    flow = AirFlow(dag_id="example_xcom")
-    res = flow.run(execution_date="2018-09-20", return_tasks=flow.tasks)
-    assert res.is_successful()
-
-    for task, state in res.result.items():
-        if task.name == "puller":
-            assert state.is_successful()
-            assert state.result is None
-        elif task.name == "push":
-            assert state.is_successful()
-            assert state.result == [1, 2, 3]
-        if task.name == "push_by_returning":
-            assert state.is_successful()
-            assert state.result == {"a": "b"}
-
-
-@pytest.mark.airflow()
-def test_example_short_circuit_operator():
-    flow = AirFlow(dag_id="example_short_circuit_operator")
-    res = flow.run(execution_date="2018-09-20", return_tasks=flow.tasks)
-    assert res.is_successful()
-
-    for task, state in res.result.items():
-        if "condition" in task.name:
-            assert isinstance(state, Success)
-        elif "false" in task.name:
-            assert isinstance(state, Skipped)
-        elif "true" in task.name:
-            assert isinstance(state, Success)
-
-
-@pytest.mark.airflow()
-def test_example_bash_operator():
-    flow = AirFlow(dag_id="example_bash_operator")
-    res = flow.run(execution_date="2018-09-20", return_tasks=flow.tasks)
-    assert res.is_successful()
-
-    for task, state in res.result.items():
-        assert isinstance(state, Success)