Add support for test_cycle from Airflow master

rambler-digital-solutions · Apr 12, 2020 · 4a764af · 4a764af
1 parent 75b1d11
commit 4a764af
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/src/airflow_declarative/builder.py b/src/airflow_declarative/builder.py
@@ -17,6 +17,7 @@
 
 from __future__ import absolute_import, division, print_function, unicode_literals
 
+from .compat import test_cycle
 from .schema import ensure_schema
 
 
@@ -89,9 +90,7 @@ def build_dag(dag_id, schema, dag_class, operator_class, sensor_class):
 
     build_flow(dict(operators, **sensors), schema.get("flow", {}))
 
-    if hasattr(dag, "test_cycle"):
-        # Since Airflow 1.10
-        dag.test_cycle()
+    test_cycle(dag)
 
     return dag
 

diff --git a/src/airflow_declarative/compat.py b/src/airflow_declarative/compat.py
@@ -19,4 +19,20 @@
 except ImportError:
     from airflow.operators import BaseOperator
 
-__all__ = ("BaseOperator", "BaseSensorOperator", "Iterable", "Mapping")
+try:
+    # Airflow master
+    from airflow.utils.dag_cycle_tester import test_cycle as _test_cycle
+except ImportError:
+    _test_cycle = None
+
+
+__all__ = ("BaseOperator", "BaseSensorOperator", "Iterable", "Mapping", "test_cycle")
+
+
+def test_cycle(dag):
+    if _test_cycle is not None:
+        # Airflow master
+        _test_cycle(dag)
+    elif hasattr(dag, "test_cycle"):
+        # Since Airflow 1.10
+        dag.test_cycle()