def test_set_dag_run_state_edge_cases(self):
# Dag does not exist
altered = set_dag_run_state(None, self.execution_dates[0])
self.assertEqual(len(altered), 0)
# Invalid execution date
altered = set_dag_run_state(self.dag1, None)
self.assertEqual(len(altered), 0)
self.assertRaises(AssertionError, set_dag_run_state, self.dag1, timedelta(microseconds=-1))
# DagRun does not exist
# This will throw AssertionError since dag.latest_execution_date does not exist
self.assertRaises(AssertionError, set_dag_run_state, self.dag1, self.execution_dates[0])
def test_set_dag_run_state_edge_cases(self):
# Dag does not exist
altered = set_dag_run_state(None, self.execution_dates[0])
self.assertEqual(len(altered), 0)
# Invalid execution date
altered = set_dag_run_state(self.dag1, None)
self.assertEqual(len(altered), 0)
self.assertRaises(AssertionError, set_dag_run_state, self.dag1, timedelta(microseconds=-1))
# DagRun does not exist
# This will throw AssertionError since dag.latest_execution_date does not exist
self.assertRaises(AssertionError, set_dag_run_state, self.dag1, self.execution_dates[0])
def test_set_mixed_dag_run_state(self):
"""
This test checks function set_dag_run_state with mixed task instance
state.
"""
date = self.execution_dates[0]
dr = self.dag1.create_dagrun(
run_id='manual__' + datetime.now().isoformat(),
state=State.FAILED,
execution_date=date,
session=self.session
)
# success task
dr.get_task_instance('runme_0').set_state(State.SUCCESS, self.session)
# skipped task
dr.get_task_instance('runme_1').set_state(State.SKIPPED, self.session)
# retry task
dr.get_task_instance('runme_2').set_state(State.UP_FOR_RETRY, self.session)
# queued task
dr.get_task_instance('also_run_this').set_state(State.QUEUED, self.session)
# running task
dr.get_task_instance('run_after_loop').set_state(State.RUNNING, self.session)
# failed task
dr.get_task_instance('run_this_last').set_state(State.FAILED, self.session)
altered = set_dag_run_state(self.dag1, date, state=State.SUCCESS, commit=True)
self.assertEqual(len(altered), len(self.dag1.tasks) - 1) # only 1 task succeeded
self.verify_dag_run_states(self.dag1, date)
def test_set_mixed_dag_run_state(self):
"""
This test checks function set_dag_run_state with mixed task instance
state.
"""
date = self.execution_dates[0]
dr = self.dag1.create_dagrun(
run_id='manual__' + datetime.now().isoformat(),
state=State.FAILED,
execution_date=date,
session=self.session
)
# success task
dr.get_task_instance('runme_0').set_state(State.SUCCESS, self.session)
# skipped task
dr.get_task_instance('runme_1').set_state(State.SKIPPED, self.session)
# retry task
dr.get_task_instance('runme_2').set_state(State.UP_FOR_RETRY, self.session)
# queued task
dr.get_task_instance('also_run_this').set_state(State.QUEUED, self.session)
# running task
dr.get_task_instance('run_after_loop').set_state(State.RUNNING, self.session)
# failed task
dr.get_task_instance('run_this_last').set_state(State.FAILED, self.session)
altered = set_dag_run_state(self.dag1, date, state=State.SUCCESS, commit=True)
self.assertEqual(len(altered), len(self.dag1.tasks) - 1) # only 1 task succeeded
self.verify_dag_run_states(self.dag1, date)
def test_set_failed_dag_run_state(self):
date = self.execution_dates[0]
dr = self.dag1.create_dagrun(
run_id='manual__' + datetime.now().isoformat(),
state=State.FAILED,
execution_date=date,
session=self.session
)
dr.get_task_instance('runme_0').set_state(State.FAILED, self.session)
altered = set_dag_run_state(self.dag1, date, state=State.SUCCESS, commit=True)
# All of the task should be altered
self.assertEqual(len(altered), len(self.dag1.tasks))
self.verify_dag_run_states(self.dag1, date)
def test_set_failed_dag_run_state(self):
date = self.execution_dates[0]
dr = self.dag1.create_dagrun(
run_id='manual__' + datetime.now().isoformat(),
state=State.FAILED,
execution_date=date,
session=self.session
)
dr.get_task_instance('runme_0').set_state(State.FAILED, self.session)
altered = set_dag_run_state(self.dag1, date, state=State.SUCCESS, commit=True)
# All of the task should be altered
self.assertEqual(len(altered), len(self.dag1.tasks))
self.verify_dag_run_states(self.dag1, date)
def test_set_running_dag_run_state(self):
date = self.execution_dates[0]
dr = self.dag1.create_dagrun(
run_id='manual__' + datetime.now().isoformat(),
state=State.RUNNING,
execution_date=date,
session=self.session
)
for ti in dr.get_task_instances(session=self.session):
ti.set_state(State.RUNNING, self.session)
altered = set_dag_run_state(self.dag1, date, state=State.SUCCESS, commit=True)
# All of the task should be altered
self.assertEqual(len(altered), len(self.dag1.tasks))
self.verify_dag_run_states(self.dag1, date)
def test_set_running_dag_run_state(self):
date = self.execution_dates[0]
dr = self.dag1.create_dagrun(
run_id='manual__' + datetime.now().isoformat(),
state=State.RUNNING,
execution_date=date,
session=self.session
)
for ti in dr.get_task_instances(session=self.session):
ti.set_state(State.RUNNING, self.session)
altered = set_dag_run_state(self.dag1, date, state=State.SUCCESS, commit=True)
# All of the task should be altered
self.assertEqual(len(altered), len(self.dag1.tasks))
self.verify_dag_run_states(self.dag1, date)
def test_set_state_without_commit(self):
date = self.execution_dates[0]
# Running dag run and task instances
dr = self.dag1.create_dagrun(
run_id='manual__' + datetime.now().isoformat(),
state=State.RUNNING,
execution_date=date,
session=self.session
)
for ti in dr.get_task_instances(session=self.session):
ti.set_state(State.RUNNING, self.session)
altered = set_dag_run_state(self.dag1, date, state=State.SUCCESS, commit=False)
# All of the task should be altered
self.assertEqual(len(altered), len(self.dag1.tasks))
# Both dag run and task instances' states should remain the same
self.verify_dag_run_states(self.dag1, date, State.RUNNING)
def test_set_state_without_commit(self):
date = self.execution_dates[0]
# Running dag run and task instances
dr = self.dag1.create_dagrun(
run_id='manual__' + datetime.now().isoformat(),
state=State.RUNNING,
execution_date=date,
session=self.session
)
for ti in dr.get_task_instances(session=self.session):
ti.set_state(State.RUNNING, self.session)
altered = set_dag_run_state(self.dag1, date, state=State.SUCCESS, commit=False)
# All of the task should be altered
self.assertEqual(len(altered), len(self.dag1.tasks))
# Both dag run and task instances' states should remain the same
self.verify_dag_run_states(self.dag1, date, State.RUNNING)
def test_set_state_with_multiple_dagruns(self):
dr1 = self.dag2.create_dagrun(run_id='manual__' +
datetime.now().isoformat(),
state=State.FAILED,
execution_date=self.execution_dates[0],
session=self.session)
dr2 = self.dag2.create_dagrun(run_id='manual__' +
datetime.now().isoformat(),
state=State.FAILED,
execution_date=self.execution_dates[1],
session=self.session)
dr3 = self.dag2.create_dagrun(run_id='manual__' +
datetime.now().isoformat(),
state=State.RUNNING,
execution_date=self.execution_dates[2],
session=self.session)
altered = set_dag_run_state(self.dag2,
self.execution_dates[1],
state=State.SUCCESS,
commit=True)
# Recursively count number of tasks in the dag
def count_dag_tasks(dag):
count = len(dag.tasks)
subdag_counts = [count_dag_tasks(subdag) for subdag in dag.subdags]
count += sum(subdag_counts)
return count
self.assertEqual(len(altered), count_dag_tasks(self.dag2))
self.verify_dag_run_states(self.dag2, self.execution_dates[1])
# Make sure other dag status are not changed
dr1 = models.DagRun.find(dag_id=self.dag2.dag_id,
execution_date=self.execution_dates[0])
dr1 = dr1[0]
self.assertEqual(dr1.get_state(), State.FAILED)
dr3 = models.DagRun.find(dag_id=self.dag2.dag_id,
execution_date=self.execution_dates[2])
dr3 = dr3[0]
self.assertEqual(dr3.get_state(), State.RUNNING)
def test_set_state_with_multiple_dagruns(self):
dr1 = self.dag2.create_dagrun(
run_id='manual__' + datetime.now().isoformat(),
state=State.FAILED,
execution_date=self.execution_dates[0],
session=self.session
)
dr2 = self.dag2.create_dagrun(
run_id='manual__' + datetime.now().isoformat(),
state=State.FAILED,
execution_date=self.execution_dates[1],
session=self.session
)
dr3 = self.dag2.create_dagrun(
run_id='manual__' + datetime.now().isoformat(),
state=State.RUNNING,
execution_date=self.execution_dates[2],
session=self.session
)
altered = set_dag_run_state(self.dag2, self.execution_dates[1],
state=State.SUCCESS, commit=True)
# Recursively count number of tasks in the dag
def count_dag_tasks(dag):
count = len(dag.tasks)
subdag_counts = [count_dag_tasks(subdag) for subdag in dag.subdags]
count += sum(subdag_counts)
return count
self.assertEqual(len(altered), count_dag_tasks(self.dag2))
self.verify_dag_run_states(self.dag2, self.execution_dates[1])
# Make sure other dag status are not changed
dr1 = models.DagRun.find(dag_id=self.dag2.dag_id, execution_date=self.execution_dates[0])
dr1 = dr1[0]
self.assertEqual(dr1.get_state(), State.FAILED)
dr3 = models.DagRun.find(dag_id=self.dag2.dag_id, execution_date=self.execution_dates[2])
dr3 = dr3[0]
self.assertEqual(dr3.get_state(), State.RUNNING)
