def test_trigger_dagrun(self):
"""Test TriggerDagRunOperator."""
task = TriggerDagRunOperator(task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, dag=self.dag)
task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
with create_session() as session:
dagruns = session.query(DagRun).filter(DagRun.dag_id == TRIGGERED_DAG_ID).all()
self.assertEqual(len(dagruns), 1)
self.assertTrue(dagruns[0].external_trigger)
def test_trigger_dagrun_with_reset_dag_run_false(self):
"""Test TriggerDagRunOperator with reset_dag_run."""
execution_date = DEFAULT_DATE
task = TriggerDagRunOperator(task_id="test_task",
trigger_dag_id=TRIGGERED_DAG_ID,
execution_date=execution_date,
reset_dag_run=False,
dag=self.dag)
task.run(start_date=execution_date, end_date=execution_date, ignore_ti_state=True)
with self.assertRaises(DagRunAlreadyExists):
task.run(start_date=execution_date, end_date=execution_date, ignore_ti_state=True)
def test_trigger_dagrun_operator_templated_conf(self):
"""Test passing a templated conf to the triggered DagRun."""
task = TriggerDagRunOperator(
task_id="test_trigger_dagrun_with_str_execution_date",
trigger_dag_id=TRIGGERED_DAG_ID,
conf={"foo": "{{ dag.dag_id }}"},
dag=self.dag,
)
task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
with create_session() as session:
dagruns = session.query(DagRun).filter(DagRun.dag_id == TRIGGERED_DAG_ID).all()
self.assertEqual(len(dagruns), 1)
self.assertTrue(dagruns[0].conf, {"foo": TEST_DAG_ID})
def test_trigger_dagrun_with_wait_for_completion_true_fail(self):
"""Test TriggerDagRunOperator with wait_for_completion but triggered dag fails."""
execution_date = DEFAULT_DATE
task = TriggerDagRunOperator(
task_id="test_task",
trigger_dag_id=TRIGGERED_DAG_ID,
execution_date=execution_date,
wait_for_completion=True,
poke_interval=10,
failed_states=[State.RUNNING],
dag=self.dag,
)
with self.assertRaises(AirflowException):
task.run(start_date=execution_date, end_date=execution_date)
def test_trigger_dagrun_with_templated_execution_date(self):
"""Test TriggerDagRunOperator with templated execution_date."""
task = TriggerDagRunOperator(
task_id="test_trigger_dagrun_with_str_execution_date",
trigger_dag_id=TRIGGERED_DAG_ID,
execution_date="{{ execution_date }}",
dag=self.dag,
)
task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
with create_session() as session:
dagruns = session.query(DagRun).filter(DagRun.dag_id == TRIGGERED_DAG_ID).all()
self.assertEqual(len(dagruns), 1)
self.assertTrue(dagruns[0].external_trigger)
self.assertEqual(dagruns[0].execution_date, DEFAULT_DATE)
def test_trigger_dagrun_with_execution_date(self):
"""Test TriggerDagRunOperator with custom execution_date."""
utc_now = timezone.utcnow()
task = TriggerDagRunOperator(
task_id="test_trigger_dagrun_with_execution_date",
trigger_dag_id=TRIGGERED_DAG_ID,
execution_date=utc_now,
dag=self.dag,
)
task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
with create_session() as session:
dagruns = session.query(DagRun).filter(DagRun.dag_id == TRIGGERED_DAG_ID).all()
self.assertEqual(len(dagruns), 1)
self.assertTrue(dagruns[0].external_trigger)
self.assertEqual(dagruns[0].execution_date, utc_now)
def test_trigger_dagrun_with_wait_for_completion_true(self):
"""Test TriggerDagRunOperator with wait_for_completion."""
execution_date = DEFAULT_DATE
task = TriggerDagRunOperator(
task_id="test_task",
trigger_dag_id=TRIGGERED_DAG_ID,
execution_date=execution_date,
wait_for_completion=True,
poke_interval=10,
allowed_states=[State.RUNNING],
dag=self.dag,
)
task.run(start_date=execution_date, end_date=execution_date)
with create_session() as session:
dagruns = session.query(DagRun).filter(DagRun.dag_id == TRIGGERED_DAG_ID).all()
self.assertEqual(len(dagruns), 1)
