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)