def test_should_render_dag_orientation(self):
orientation = "TB"
dag = DAG(dag_id="DAG_ID", orientation=orientation)
task_1 = BashOperator(dag=dag, start_date=START_DATE, task_id="first", bash_command="echo 1")
task_2 = BashOperator(dag=dag, start_date=START_DATE, task_id="second", bash_command="echo 1")
task_3 = PythonOperator(
dag=dag, start_date=START_DATE, task_id="third", python_callable=mock.MagicMock()
)
task_1 >> task_2
task_1 >> task_3
tis = [
TaskInstance(task_1, execution_date=START_DATE, state=State.SCHEDULED),
TaskInstance(task_2, execution_date=START_DATE, state=State.SUCCESS),
TaskInstance(task_3, execution_date=START_DATE, state=State.RUNNING),
]
dot = dot_renderer.render_dag(dag, tis=tis)
source = dot.source
# Should render DAG title with orientation
assert "label=DAG_ID" in source
assert f'label=DAG_ID labelloc=t rankdir={orientation}' in source
# Change orientation
orientation = "LR"
dag = DAG(dag_id="DAG_ID", orientation=orientation)
dot = dot_renderer.render_dag(dag, tis=tis)
source = dot.source
# Should render DAG title with orientation
assert "label=DAG_ID" in source
assert f'label=DAG_ID labelloc=t rankdir={orientation}' in source
def dag_show(args):
"""Displays DAG or saves it's graphic representation to the file"""
dag = get_dag(args.subdir, args.dag_id)
dot = render_dag(dag)
if args.save:
filename, _, fileformat = args.save.rpartition('.')
dot.render(filename=filename, format=fileformat, cleanup=True)
print("File {} saved".format(args.save))
elif args.imgcat:
data = dot.pipe(format='png')
try:
proc = subprocess.Popen("imgcat", stdout=subprocess.PIPE, stdin=subprocess.PIPE)
except OSError as e:
if e.errno == errno.ENOENT:
raise AirflowException(
"Failed to execute. Make sure the imgcat executables are on your systems \'PATH\'"
)
else:
raise
out, err = proc.communicate(data)
if out:
print(out.decode('utf-8'))
if err:
print(err.decode('utf-8'))
else:
print(dot.source)
def dag_test(args, session=None):
"""Execute one single DagRun for a given DAG and execution date, using the DebugExecutor."""
dag = get_dag(subdir=args.subdir, dag_id=args.dag_id)
dag.clear(start_date=args.execution_date,
end_date=args.execution_date,
dag_run_state=State.NONE)
try:
dag.run(executor=DebugExecutor(),
start_date=args.execution_date,
end_date=args.execution_date)
except BackfillUnfinished as e:
print(str(e))
show_dagrun = args.show_dagrun
imgcat = args.imgcat_dagrun
filename = args.save_dagrun
if show_dagrun or imgcat or filename:
tis = session.query(TaskInstance).filter(
TaskInstance.dag_id == args.dag_id,
TaskInstance.execution_date == args.execution_date,
).all()
dot_graph = render_dag(dag, tis=tis)
print()
if filename:
_save_dot_to_file(dot_graph, filename)
if imgcat:
_display_dot_via_imgcat(dot_graph)
if show_dagrun:
print(dot_graph.source)
def test_should_render_dag(self):
dag = DAG(dag_id="DAG_ID")
task_1 = BashOperator(dag=dag,
start_date=START_DATE,
task_id="first",
bash_command="echo 1")
task_2 = BashOperator(dag=dag,
start_date=START_DATE,
task_id="second",
bash_command="echo 1")
task_3 = PythonOperator(dag=dag,
start_date=START_DATE,
task_id="third",
python_callable=mock.MagicMock())
task_1 >> task_2
task_1 >> task_3
dot = dot_renderer.render_dag(dag)
source = dot.source
# Should render DAG title
self.assertIn("label=DAG_ID", source)
self.assertIn("first", source)
self.assertIn("second", source)
self.assertIn("third", source)
self.assertIn("first -> second", source)
self.assertIn("first -> third", source)
self.assertIn('fillcolor="#f0ede4"', source)
self.assertIn('fillcolor="#f0ede4"', source)
def test_should_render_dag_with_task_instances(self):
dag = DAG(dag_id="DAG_ID")
task_1 = BashOperator(dag=dag, start_date=START_DATE, task_id="first", bash_command="echo 1")
task_2 = BashOperator(dag=dag, start_date=START_DATE, task_id="second", bash_command="echo 1")
task_3 = PythonOperator(
dag=dag, start_date=START_DATE, task_id="third", python_callable=mock.MagicMock()
)
task_1 >> task_2
task_1 >> task_3
tis = [
TaskInstance(task_1, execution_date=START_DATE, state=State.SCHEDULED),
TaskInstance(task_2, execution_date=START_DATE, state=State.SUCCESS),
TaskInstance(task_3, execution_date=START_DATE, state=State.RUNNING),
]
dot = dot_renderer.render_dag(dag, tis=tis)
source = dot.source
# Should render DAG title
assert "label=DAG_ID" in source
assert (
'first [color=black fillcolor=tan label=first shape=rectangle style="filled,rounded"]' in source
)
assert (
'second [color=white fillcolor=green label=second shape=rectangle style="filled,rounded"]'
in source
)
assert (
'third [color=black fillcolor=lime label=third shape=rectangle style="filled,rounded"]' in source
)
def test_disable_task(self):
dag_creator = DagCreator(self.task_graph._graph, with_data_nodes=True)
dags = dag_creator.traverse_graph()
self.assertEqual(len(dags), 3)
test_spark_dag = dags['test_spark']
dot = render_dag(test_spark_dag)
self.assertEqual(dot.source, self.dot_test_spark_deactivate)
def test_dag_creator_external_sensor(self):
dag_creator = DagCreator(self.task_graph._graph)
dags = dag_creator.traverse_graph()
self.assertEqual(len(dags), 3)
test_external_sensor_dag = dags['test_external_sensor']
dot = render_dag(test_external_sensor_dag)
self.assertEqual(dot.source, self.dot_test_external_sensor)
def test_dag_creator_with_dataset(self):
dag_creator = DagCreator(self.task_graph._graph, with_data_nodes=True)
dags = dag_creator.traverse_graph()
self.assertEqual(len(dags), 3)
test_batch_dag = dags['test_batch']
dot = render_dag(test_batch_dag)
self.assertEqual(dot.source, self.dot_test_batch_graph_with_dataset)
def dag_show(args):
"""Displays DAG or saves it's graphic representation to the file"""
dag = get_dag(args.subdir, args.dag_id)
dot = render_dag(dag)
filename = args.save
imgcat = args.imgcat
if filename and imgcat:
raise SystemExit(
"Option --save and --imgcat are mutually exclusive. "
"Please remove one option to execute the command.", )
elif filename:
_save_dot_to_file(dot, filename)
elif imgcat:
_display_dot_via_imgcat(dot)
else:
print(dot.source)
def dag_test(args, session=None):
"""Execute one single DagRun for a given DAG and execution date, using the DebugExecutor."""
dag = get_dag(subdir=args.subdir, dag_id=args.dag_id)
dag.clear(start_date=args.execution_date, end_date=args.execution_date, dag_run_state=False)
try:
dag.run(
executor=DebugExecutor(),
start_date=args.execution_date,
end_date=args.execution_date,
# Always run the DAG at least once even if no logical runs are
# available. This does not make a lot of sense, but Airflow has
# been doing this prior to 2.2 so we keep compatibility.
run_at_least_once=True,
)
except BackfillUnfinished as e:
print(str(e))
show_dagrun = args.show_dagrun
imgcat = args.imgcat_dagrun
filename = args.save_dagrun
if show_dagrun or imgcat or filename:
tis = (
session.query(TaskInstance)
.filter(
TaskInstance.dag_id == args.dag_id,
TaskInstance.execution_date == args.execution_date,
)
.all()
)
dot_graph = render_dag(dag, tis=tis)
print()
if filename:
_save_dot_to_file(dot_graph, filename)
if imgcat:
_display_dot_via_imgcat(dot_graph)
if show_dagrun:
print(dot_graph.source)
def test_render_task_group(self):
with DAG(dag_id="example_task_group", start_date=START_DATE) as dag:
start = DummyOperator(task_id="start")
with TaskGroup("section_1", tooltip="Tasks for section_1") as section_1:
task_1 = DummyOperator(task_id="task_1")
task_2 = BashOperator(task_id="task_2", bash_command='echo 1')
task_3 = DummyOperator(task_id="task_3")
task_1 >> [task_2, task_3]
with TaskGroup("section_2", tooltip="Tasks for section_2") as section_2:
task_1 = DummyOperator(task_id="task_1")
with TaskGroup("inner_section_2", tooltip="Tasks for inner_section2"):
task_2 = BashOperator(task_id="task_2", bash_command='echo 1')
task_3 = DummyOperator(task_id="task_3")
task_4 = DummyOperator(task_id="task_4")
[task_2, task_3] >> task_4
end = DummyOperator(task_id='end')
start >> section_1 >> section_2 >> end
dot = dot_renderer.render_dag(dag)
assert dot.source == '\n'.join(
[
'digraph example_task_group {',
'\tgraph [label=example_task_group labelloc=t rankdir=LR]',
'\tend [color="#000000" fillcolor="#e8f7e4" label=end shape=rectangle '
'style="filled,rounded"]',
'\tsubgraph cluster_section_1 {',
'\t\tcolor="#000000" fillcolor="#6495ed7f" label=section_1 shape=rectangle style=filled',
'\t\t"section_1.upstream_join_id" [color="#000000" fillcolor=CornflowerBlue height=0.2 '
'label="" shape=circle style="filled,rounded" width=0.2]',
'\t\t"section_1.downstream_join_id" [color="#000000" fillcolor=CornflowerBlue height=0.2 '
'label="" shape=circle style="filled,rounded" width=0.2]',
'\t\t"section_1.task_1" [color="#000000" fillcolor="#e8f7e4" label=task_1 shape=rectangle '
'style="filled,rounded"]',
'\t\t"section_1.task_2" [color="#000000" fillcolor="#f0ede4" label=task_2 shape=rectangle '
'style="filled,rounded"]',
'\t\t"section_1.task_3" [color="#000000" fillcolor="#e8f7e4" label=task_3 shape=rectangle '
'style="filled,rounded"]',
'\t}',
'\tsubgraph cluster_section_2 {',
'\t\tcolor="#000000" fillcolor="#6495ed7f" label=section_2 shape=rectangle style=filled',
'\t\t"section_2.upstream_join_id" [color="#000000" fillcolor=CornflowerBlue height=0.2 '
'label="" shape=circle style="filled,rounded" width=0.2]',
'\t\t"section_2.downstream_join_id" [color="#000000" fillcolor=CornflowerBlue height=0.2 '
'label="" shape=circle style="filled,rounded" width=0.2]',
'\t\tsubgraph "cluster_section_2.inner_section_2" {',
'\t\t\tcolor="#000000" fillcolor="#6495ed7f" label=inner_section_2 shape=rectangle '
'style=filled',
'\t\t\t"section_2.inner_section_2.task_2" [color="#000000" fillcolor="#f0ede4" label=task_2 '
'shape=rectangle style="filled,rounded"]',
'\t\t\t"section_2.inner_section_2.task_3" [color="#000000" fillcolor="#e8f7e4" label=task_3 '
'shape=rectangle style="filled,rounded"]',
'\t\t\t"section_2.inner_section_2.task_4" [color="#000000" fillcolor="#e8f7e4" label=task_4 '
'shape=rectangle style="filled,rounded"]',
'\t\t}',
'\t\t"section_2.task_1" [color="#000000" fillcolor="#e8f7e4" label=task_1 shape=rectangle '
'style="filled,rounded"]',
'\t}',
'\tstart [color="#000000" fillcolor="#e8f7e4" label=start shape=rectangle '
'style="filled,rounded"]',
'\t"section_1.downstream_join_id" -> "section_2.upstream_join_id"',
'\t"section_1.task_1" -> "section_1.task_2"',
'\t"section_1.task_1" -> "section_1.task_3"',
'\t"section_1.task_2" -> "section_1.downstream_join_id"',
'\t"section_1.task_3" -> "section_1.downstream_join_id"',
'\t"section_1.upstream_join_id" -> "section_1.task_1"',
'\t"section_2.downstream_join_id" -> end',
'\t"section_2.inner_section_2.task_2" -> "section_2.inner_section_2.task_4"',
'\t"section_2.inner_section_2.task_3" -> "section_2.inner_section_2.task_4"',
'\t"section_2.inner_section_2.task_4" -> "section_2.downstream_join_id"',
'\t"section_2.task_1" -> "section_2.downstream_join_id"',
'\t"section_2.upstream_join_id" -> "section_2.inner_section_2.task_2"',
'\t"section_2.upstream_join_id" -> "section_2.inner_section_2.task_3"',
'\t"section_2.upstream_join_id" -> "section_2.task_1"',
'\tstart -> "section_1.upstream_join_id"',
'}',
]
)
