def test_show_dag_print(self):
temp_stdout = io.StringIO()
with contextlib.redirect_stdout(temp_stdout):
cli.show_dag(self.parser.parse_args([
'dags', 'show', 'example_bash_operator']))
out = temp_stdout.getvalue()
self.assertIn("label=example_bash_operator", out)
self.assertIn("graph [label=example_bash_operator labelloc=t rankdir=LR]", out)
self.assertIn("runme_2 -> run_after_loop", 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")
temp_stdout = ByteableIO() if PY2 else io.StringIO()
with redirect_stdout(temp_stdout):
cli.show_dag(self.parser.parse_args([
'show_dag', 'example_bash_operator', '--imgcat']
))
mock_render_dag.return_value.pipe.assert_called_once_with(format='png')
mock_popen.return_value.communicate.assert_called_once_with(b'DOT_DATA')
def test_show_dag_dave(self, mock_render_dag):
temp_stdout = io.StringIO()
with contextlib.redirect_stdout(temp_stdout):
cli.show_dag(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'
)
self.assertIn("File awesome.png saved", 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")
temp_stdout = io.StringIO()
with contextlib.redirect_stdout(temp_stdout):
cli.show_dag(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)
