test: add tests for plot in end, plot in middle, plot twice

tests/unit/flow/test_flow_visualization.py

@@ -38,4 +38,25 @@ def test_visualization_plot_in_middle(tmpdir):
     with Image.open(os.path.join(cur_dir, 'flow_original_middle.jpg')) as flow_original:
         with Image.open(tmpfile) as flow_created:
             assert flow_original.size == flow_created.size
-            np.testing.assert_array_almost_equal(flow_original, flow_created)
+            np.testing.assert_array_almost_equal(flow_original, flow_created)
+
+
+def test_visualization_plot_twice(tmpdir):
+    cur_dir = os.path.dirname(os.path.abspath(__file__))
+    tmpfile_middle = os.path.join(tmpdir, 'flow_test_middle.jpg')
+    tmpfile_end = os.path.join(tmpdir, 'flow_test_end.jpg')
+
+    flow = (Flow().add(name='pod_a')
+            .plot(output=tmpfile_middle)
+            .add(name='pod_b', needs='gateway')
+            .join(needs=['pod_a', 'pod_b']).plot(output=tmpfile_end))
+
+    with Image.open(os.path.join(cur_dir, 'flow_original_middle.jpg')) as flow_original_middle:
+        with Image.open(tmpfile_middle) as flow_created_middle:
+            assert flow_original_middle.size == flow_created_middle.size
+            np.testing.assert_array_almost_equal(flow_original_middle, flow_created_middle)
+
+    with Image.open(os.path.join(cur_dir, 'flow_original_end.jpg')) as flow_original_end:
+        with Image.open(tmpfile_end) as flow_created_end:
+            assert flow_original_end.size == flow_created_end.size
+            np.testing.assert_array_almost_equal(flow_original_end, flow_created_end)