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})
