def test_set_dag_run_state_edge_cases(self):
# Dag does not exist
altered = set_dag_run_state_to_success(None, self.execution_dates[0])
self.assertEqual(len(altered), 0)
altered = set_dag_run_state_to_failed(None, self.execution_dates[0])
self.assertEqual(len(altered), 0)
altered = set_dag_run_state_to_running(None, self.execution_dates[0])
self.assertEqual(len(altered), 0)
# Invalid execution date
altered = set_dag_run_state_to_success(self.dag1, None)
self.assertEqual(len(altered), 0)
altered = set_dag_run_state_to_failed(self.dag1, None)
self.assertEqual(len(altered), 0)
altered = set_dag_run_state_to_running(self.dag1, None)
self.assertEqual(len(altered), 0)
# This will throw AssertionError since dag.latest_execution_date
# need to be 0 does not exist.
self.assertRaises(AssertionError, set_dag_run_state_to_success, self.dag1,
timezone.make_naive(self.execution_dates[0]))
# altered = set_dag_run_state_to_success(self.dag1, self.execution_dates[0])
# DagRun does not exist
# This will throw AssertionError since dag.latest_execution_date does not exist
self.assertRaises(AssertionError, set_dag_run_state_to_success,
self.dag1, self.execution_dates[0])
def test_set_dag_run_state_edge_cases(self):
# Dag does not exist
altered = set_dag_run_state_to_success(None, self.execution_dates[0])
assert len(altered) == 0
altered = set_dag_run_state_to_failed(None, self.execution_dates[0])
assert len(altered) == 0
altered = set_dag_run_state_to_running(None, self.execution_dates[0])
assert len(altered) == 0
# Invalid execution date
altered = set_dag_run_state_to_success(self.dag1, None)
assert len(altered) == 0
altered = set_dag_run_state_to_failed(self.dag1, None)
assert len(altered) == 0
altered = set_dag_run_state_to_running(self.dag1, None)
assert len(altered) == 0
# This will throw ValueError since dag.latest_execution_date
# need to be 0 does not exist.
with pytest.raises(ValueError):
set_dag_run_state_to_success(
self.dag2, timezone.make_naive(self.execution_dates[0]))
# altered = set_dag_run_state_to_success(self.dag1, self.execution_dates[0])
# DagRun does not exist
# This will throw ValueError since dag.latest_execution_date does not exist
with pytest.raises(ValueError):
set_dag_run_state_to_success(self.dag2, self.execution_dates[0])
def test_set_dag_run_state_edge_cases(self):
# Dag does not exist
altered = set_dag_run_state_to_success(None, self.execution_dates[0])
self.assertEqual(len(altered), 0)
altered = set_dag_run_state_to_failed(None, self.execution_dates[0])
self.assertEqual(len(altered), 0)
altered = set_dag_run_state_to_running(None, self.execution_dates[0])
self.assertEqual(len(altered), 0)
# Invalid execution date
altered = set_dag_run_state_to_success(self.dag1, None)
self.assertEqual(len(altered), 0)
altered = set_dag_run_state_to_failed(self.dag1, None)
self.assertEqual(len(altered), 0)
altered = set_dag_run_state_to_running(self.dag1, None)
self.assertEqual(len(altered), 0)
# This will throw ValueError since dag.latest_execution_date
# need to be 0 does not exist.
self.assertRaises(ValueError, set_dag_run_state_to_success, self.dag2,
timezone.make_naive(self.execution_dates[0]))
# altered = set_dag_run_state_to_success(self.dag1, self.execution_dates[0])
# DagRun does not exist
# This will throw ValueError since dag.latest_execution_date does not exist
self.assertRaises(ValueError, set_dag_run_state_to_success, self.dag2,
self.execution_dates[0])
def test_set_state_without_commit(self):
date = self.execution_dates[0]
dr = self._create_test_dag_run(State.RUNNING, date)
self._set_default_task_instance_states(dr)
will_be_altered = set_dag_run_state_to_running(self.dag1, date, commit=False)
# None of the tasks will be altered.
self.assertEqual(len(will_be_altered), 0)
self._verify_dag_run_state(self.dag1, date, State.RUNNING)
self._verify_task_instance_states_remain_default(dr)
will_be_altered = set_dag_run_state_to_failed(self.dag1, date, commit=False)
# Only the running task will be altered.
self.assertEqual(len(will_be_altered), 1)
self._verify_dag_run_state(self.dag1, date, State.RUNNING)
self._verify_task_instance_states_remain_default(dr)
will_be_altered = set_dag_run_state_to_success(self.dag1, date, commit=False)
# All except the SUCCESS task should be altered.
self.assertEqual(len(will_be_altered), 5)
self._verify_dag_run_state(self.dag1, date, State.RUNNING)
self._verify_task_instance_states_remain_default(dr)
def test_set_state_without_commit(self):
date = self.execution_dates[0]
dr = self._create_test_dag_run(State.RUNNING, date)
self._set_default_task_instance_states(dr)
will_be_altered = set_dag_run_state_to_running(self.dag1,
date,
commit=False)
# None of the tasks will be altered.
self.assertEqual(len(will_be_altered), 0)
self._verify_dag_run_state(self.dag1, date, State.RUNNING)
self._verify_task_instance_states_remain_default(dr)
will_be_altered = set_dag_run_state_to_failed(self.dag1,
date,
commit=False)
# Only the running task will be altered.
self.assertEqual(len(will_be_altered), 1)
self._verify_dag_run_state(self.dag1, date, State.RUNNING)
self._verify_task_instance_states_remain_default(dr)
will_be_altered = set_dag_run_state_to_success(self.dag1,
date,
commit=False)
# All except the SUCCESS task should be altered.
self.assertEqual(len(will_be_altered), 5)
self._verify_dag_run_state(self.dag1, date, State.RUNNING)
self._verify_task_instance_states_remain_default(dr)
def test_set_state_without_commit(self):
date = self.execution_dates[0]
dr = self._create_test_dag_run(State.RUNNING, date)
self._set_default_task_instance_states(dr)
will_be_altered = set_dag_run_state_to_running(self.dag1,
date,
commit=False)
# None of the tasks will be altered.
assert len(will_be_altered) == 0
self._verify_dag_run_state(self.dag1, date, State.RUNNING)
self._verify_task_instance_states_remain_default(dr)
will_be_altered = set_dag_run_state_to_failed(self.dag1,
date,
commit=False)
# Only the running task should be altered.
expected = self._get_num_tasks_with_starting_state(State.RUNNING,
inclusion=True)
assert len(will_be_altered) == expected
self._verify_dag_run_state(self.dag1, date, State.RUNNING)
self._verify_task_instance_states_remain_default(dr)
will_be_altered = set_dag_run_state_to_success(self.dag1,
date,
commit=False)
# All except the SUCCESS task should be altered.
expected = self._get_num_tasks_with_starting_state(State.SUCCESS,
inclusion=False)
assert len(will_be_altered) == expected
self._verify_dag_run_state(self.dag1, date, State.RUNNING)
self._verify_task_instance_states_remain_default(dr)
def test_set_failed_dag_run_to_running(self):
date = self.execution_dates[0]
dr = self._create_test_dag_run(State.SUCCESS, date)
self._set_default_task_instance_states(dr)
altered = set_dag_run_state_to_running(self.dag1, date, commit=True)
# None of the tasks should be altered.
self.assertEqual(len(altered), 0)
self._verify_dag_run_state(self.dag1, date, State.RUNNING)
self._verify_task_instance_states_remain_default(dr)
def test_set_failed_dag_run_to_running(self):
date = self.execution_dates[0]
dr = self._create_test_dag_run(State.SUCCESS, date)
middle_time = timezone.utcnow()
self._set_default_task_instance_states(dr)
altered = set_dag_run_state_to_running(self.dag1, date, commit=True)
# None of the tasks should be altered.
self.assertEqual(len(altered), 0)
self._verify_dag_run_state(self.dag1, date, State.RUNNING)
self._verify_task_instance_states_remain_default(dr)
self._verify_dag_run_dates(self.dag1, date, State.RUNNING, middle_time)
def test_set_success_dag_run_to_running(self):
date = self.execution_dates[0]
dr = self._create_test_dag_run(State.SUCCESS, date)
middle_time = timezone.utcnow()
self._set_default_task_instance_states(dr)
altered = set_dag_run_state_to_running(self.dag1, date, commit=True)
# None of the tasks should be altered, but only the dag object should be changed
self.assertEqual(len(altered), 0)
self._verify_dag_run_state(self.dag1, date, State.RUNNING)
self._verify_task_instance_states_remain_default(dr)
self._verify_dag_run_dates(self.dag1, date, State.RUNNING, middle_time)
def test_set_failed_dag_run_to_running(self):
date = self.execution_dates[0]
dr = self._create_test_dag_run(State.SUCCESS, date)
middle_time = timezone.utcnow()
self._set_default_task_instance_states(dr)
altered = set_dag_run_state_to_running(self.dag1, date, commit=True)
# None of the tasks should be altered, since we've only altered the DAG itself
assert len(altered) == 0
self._verify_dag_run_state(self.dag1, date, State.RUNNING)
self._verify_task_instance_states_remain_default(dr)
self._verify_dag_run_dates(self.dag1, date, State.RUNNING, middle_time)
