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_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,
        )

        self.assertEqual(3, len(subdag_task._get_dagrun.mock_calls))
Exemplo n.º 3
0
    def test_execute_create_dagrun_wait_until_success(self):
        """
        When SubDagOperator executes, it creates a DagRun if there is no existing one
        and wait until the DagRun succeeds.
        """
        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_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_id="scheduled__{}".format(DEFAULT_DATE.isoformat()),
            execution_date=DEFAULT_DATE,
            state=State.RUNNING,
            external_trigger=True,
        )

        self.assertEqual(3, len(subdag_task._get_dagrun.mock_calls))
    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()
        self.assertEqual(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 self.assertRaises(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})