def test_execute_create_dagrun_with_conf(self): """ When SubDagOperator executes, it creates a DagRun if there is no existing one and wait until the DagRun succeeds. """ conf = {"key": "value"} dag = DAG('parent', default_args=default_args) subdag = DAG('parent.test', default_args=default_args) subdag_task = SubDagOperator(task_id='test', subdag=subdag, dag=dag, poke_interval=1, conf=conf) subdag.create_dagrun = Mock() subdag.create_dagrun.return_value = self.dag_run_running subdag_task._get_dagrun = Mock() subdag_task._get_dagrun.side_effect = [ None, self.dag_run_success, self.dag_run_success ] subdag_task.pre_execute(context={'execution_date': DEFAULT_DATE}) subdag_task.execute(context={'execution_date': DEFAULT_DATE}) subdag_task.post_execute(context={'execution_date': DEFAULT_DATE}) subdag.create_dagrun.assert_called_once_with( run_type=DagRunType.SCHEDULED, execution_date=DEFAULT_DATE, conf=conf, state=State.RUNNING, external_trigger=True, ) assert 3 == len(subdag_task._get_dagrun.mock_calls)
def test_subdag_with_propagate_skipped_state(self, propagate_option, states, skip_parent, mock_get_task_instance, mock_skip): """ Tests that skipped state of leaf tasks propagates to the parent dag. Note that the skipped state propagation only takes affect when the dagrun's state is SUCCESS. """ dag = DAG('parent', default_args=default_args) subdag = DAG('parent.test', default_args=default_args) subdag_task = SubDagOperator(task_id='test', subdag=subdag, dag=dag, poke_interval=1, propagate_skipped_state=propagate_option) dummy_subdag_tasks = [ DummyOperator(task_id=f'dummy_subdag_{i}', dag=subdag) for i in range(len(states)) ] dummy_dag_task = DummyOperator(task_id='dummy_dag', dag=dag) subdag_task >> dummy_dag_task subdag_task._get_dagrun = Mock() subdag_task._get_dagrun.return_value = self.dag_run_success mock_get_task_instance.side_effect = [ TaskInstance(task=task, execution_date=DEFAULT_DATE, state=state) for task, state in zip(dummy_subdag_tasks, states) ] context = { 'execution_date': DEFAULT_DATE, 'dag_run': DagRun(), 'task': subdag_task } subdag_task.post_execute(context) if skip_parent: mock_skip.assert_called_once_with(context['dag_run'], context['execution_date'], [dummy_dag_task]) else: mock_skip.assert_not_called()
def test_execute_skip_if_dagrun_success(self): """ When there is an existing DagRun in SUCCESS state, skip the execution. """ dag = DAG('parent', default_args=default_args) subdag = DAG('parent.test', default_args=default_args) subdag.create_dagrun = Mock() subdag_task = SubDagOperator(task_id='test', subdag=subdag, dag=dag, poke_interval=1) subdag_task._get_dagrun = Mock() subdag_task._get_dagrun.return_value = self.dag_run_success subdag_task.pre_execute(context={'execution_date': DEFAULT_DATE}) subdag_task.execute(context={'execution_date': DEFAULT_DATE}) subdag_task.post_execute(context={'execution_date': DEFAULT_DATE}) subdag.create_dagrun.assert_not_called() assert 3 == len(subdag_task._get_dagrun.mock_calls)
def test_execute_dagrun_failed(self): """ When the DagRun failed during the execution, it raises an Airflow Exception. """ dag = DAG('parent', default_args=default_args) subdag = DAG('parent.test', default_args=default_args) subdag_task = SubDagOperator(task_id='test', subdag=subdag, dag=dag, poke_interval=1) subdag.create_dagrun = Mock() subdag.create_dagrun.return_value = self.dag_run_running subdag_task._get_dagrun = Mock() subdag_task._get_dagrun.side_effect = [ None, self.dag_run_failed, self.dag_run_failed ] with pytest.raises(AirflowException): subdag_task.pre_execute(context={'execution_date': DEFAULT_DATE}) subdag_task.execute(context={'execution_date': DEFAULT_DATE}) subdag_task.post_execute(context={'execution_date': DEFAULT_DATE})