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()
def test_cycle_empty(self): # test empty dag = DAG('dag', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'}) self.assertFalse(_test_cycle(dag))
def test_cycle_single_task(self): # test single task dag = DAG('dag', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'}) with dag: DummyOperator(task_id='A') self.assertFalse(_test_cycle(dag))
def test_cycle_loop(self): # test self loop dag = DAG('dag', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'}) # A -> A with dag: op1 = DummyOperator(task_id='A') op1.set_downstream(op1) with self.assertRaises(AirflowDagCycleException): self.assertFalse(_test_cycle(dag))
def test_cycle_large_loop(self): # large loop dag = DAG('dag', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'}) # A -> B -> C -> D -> E -> A with dag: start = DummyOperator(task_id='start') current = start for i in range(10000): next_task = DummyOperator(task_id=f'task_{i}') current.set_downstream(next_task) current = next_task current.set_downstream(start) with self.assertRaises(AirflowDagCycleException): self.assertFalse(_test_cycle(dag))
def test_cycle_downstream_loop(self): # test downstream self loop dag = DAG('dag', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'}) # A -> B -> C -> D -> E -> E with dag: op1 = DummyOperator(task_id='A') op2 = DummyOperator(task_id='B') op3 = DummyOperator(task_id='C') op4 = DummyOperator(task_id='D') op5 = DummyOperator(task_id='E') op1.set_downstream(op2) op2.set_downstream(op3) op3.set_downstream(op4) op4.set_downstream(op5) op5.set_downstream(op5) with self.assertRaises(AirflowDagCycleException): self.assertFalse(_test_cycle(dag))
def test_cycle_no_cycle(self): # test no cycle dag = DAG('dag', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'}) # A -> B -> C # B -> D # E -> F with dag: op1 = DummyOperator(task_id='A') op2 = DummyOperator(task_id='B') op3 = DummyOperator(task_id='C') op4 = DummyOperator(task_id='D') op5 = DummyOperator(task_id='E') op6 = DummyOperator(task_id='F') op1.set_downstream(op2) op2.set_downstream(op3) op2.set_downstream(op4) op5.set_downstream(op6) self.assertFalse(_test_cycle(dag))
def test_cycle_arbitrary_loop(self): # test arbitrary loop dag = DAG('dag', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'}) # E-> A -> B -> F -> A # -> C -> F with dag: op1 = DummyOperator(task_id='A') op2 = DummyOperator(task_id='B') op3 = DummyOperator(task_id='C') op4 = DummyOperator(task_id='E') op5 = DummyOperator(task_id='F') op1.set_downstream(op2) op1.set_downstream(op3) op4.set_downstream(op1) op3.set_downstream(op5) op2.set_downstream(op5) op5.set_downstream(op1) with pytest.raises(AirflowDagCycleException): assert not _test_cycle(dag)