def test_show_dag_print(self):
with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
dag_command.dag_show(self.parser.parse_args(['dags', 'show', 'example_bash_operator']))
out = temp_stdout.getvalue()
assert "label=example_bash_operator" in out
assert "graph [label=example_bash_operator labelloc=t rankdir=LR]" in out
assert "runme_2 -> run_after_loop" in out
def test_show_dag_dave(self, mock_render_dag):
with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
dag_command.dag_show(
self.parser.parse_args(['dags', 'show', 'example_bash_operator', '--save', 'awesome.png'])
)
out = temp_stdout.getvalue()
mock_render_dag.return_value.render.assert_called_once_with(
cleanup=True, filename='awesome', format='png'
)
assert "File awesome.png saved" in out
def test_show_dag_imgcat(self, mock_render_dag, mock_popen):
mock_render_dag.return_value.pipe.return_value = b"DOT_DATA"
mock_popen.return_value.communicate.return_value = (b"OUT", b"ERR")
with contextlib.redirect_stdout(io.StringIO()) as temp_stdout:
dag_command.dag_show(self.parser.parse_args([
'dags', 'show', 'example_bash_operator', '--imgcat']
))
out = temp_stdout.getvalue()
mock_render_dag.return_value.pipe.assert_called_once_with(format='png')
mock_popen.return_value.communicate.assert_called_once_with(b'DOT_DATA')
self.assertIn("OUT", out)
self.assertIn("ERR", out)