def test_task_states_for_dag_run(self):
dag2 = DagBag().dags['example_python_operator']
task2 = dag2.get_task(task_id='print_the_context')
defaut_date2 = timezone.make_aware(datetime(2016, 1, 9))
dag2.clear()
ti2 = TaskInstance(task2, defaut_date2)
ti2.set_state(State.SUCCESS)
ti_start = ti2.start_date
ti_end = ti2.end_date
with redirect_stdout(io.StringIO()) as stdout:
task_command.task_states_for_dag_run(
self.parser.parse_args([
'tasks', 'states-for-dag-run', 'example_python_operator',
defaut_date2.isoformat()
]))
actual_out = stdout.getvalue()
formatted_rows = [
('example_python_operator', '2016-01-09 00:00:00+00:00',
'print_the_context', 'success', ti_start, ti_end)
]
expected = tabulate(
formatted_rows,
['dag', 'exec_date', 'task', 'state', 'start_date', 'end_date'],
tablefmt="plain")
# Check that prints, and log messages, are shown
self.assertIn(expected.replace("\n", ""), actual_out.replace("\n", ""))
def test_task_states_for_dag_run(self):
dag2 = DagBag().dags['example_python_operator']
task2 = dag2.get_task(task_id='print_the_context')
default_date2 = timezone.make_aware(datetime(2016, 1, 9))
dag2.clear()
ti2 = TaskInstance(task2, default_date2)
ti2.set_state(State.SUCCESS)
ti_start = ti2.start_date
ti_end = ti2.end_date
with redirect_stdout(io.StringIO()) as stdout:
task_command.task_states_for_dag_run(
self.parser.parse_args([
'tasks',
'states-for-dag-run',
'example_python_operator',
default_date2.isoformat(),
'--output',
"json",
]))
actual_out = json.loads(stdout.getvalue())
assert len(actual_out) == 1
assert actual_out[0] == {
'dag_id': 'example_python_operator',
'execution_date': '2016-01-09T00:00:00+00:00',
'task_id': 'print_the_context',
'state': 'success',
'start_date': ti_start.isoformat(),
'end_date': ti_end.isoformat(),
}
def clear_subdag_callback(context):
dag_id = "{}.{}".format(context['dag'].dag_id, context['ti'].task_id)
execution_date = context['execution_date']
sdag = DagBag().get_dag(dag_id)
sdag.clear(start_date=execution_date,
end_date=execution_date,
only_failed=False,
only_running=False,
confirm_prompt=False,
include_subdags=False)