add additional tests

wntr/tests/test_graphics.py

@@ -117,6 +117,52 @@ def test_plot_network5(self):
         plt.close()
 
         self.assertTrue(isfile(filename))
+
+    def test_plot_network6(self):
+        # pumps/valves
+        filename = abspath(join(testdir, "plot_network6.png"))
+        if isfile(filename):
+            os.remove(filename)
+
+        inp_file = join(ex_datadir, "Net6.inp")
+        wn = wntr.network.WaterNetworkModel(inp_file)
+
+        # verify that direction points away from start node
+        start_nodes = []
+        for link_name in wn.pump_name_list+wn.valve_name_list:
+            link = wn.get_link(link_name)
+            start_nodes.append(link.start_node_name)
+
+        # pump=0, valve=1
+        link_type = pd.Series(0, index=wn.pump_name_list+wn.valve_name_list)
+        link_type[wn.valve_name_list] = 1
+
+        wntr.graphics.plot_network(
+            wn, node_attribute=start_nodes, link_attribute=link_type,
+            show_pump_direction=True, show_valve_direction=True
+        )
+        plt.savefig(filename, format="png")
+        plt.close()
+
+        self.assertTrue(isfile(filename))
+
+    def test_plot_network7(self):
+        # legend
+        filename = abspath(join(testdir, "plot_network7.png"))
+        if isfile(filename):
+            os.remove(filename)
+
+        inp_file = join(ex_datadir, "Net6.inp")
+        wn = wntr.network.WaterNetworkModel(inp_file)
+
+        wntr.graphics.plot_network(
+            wn, node_attribute="elevation", link_attribute="diameter", 
+            add_colorbar=True, legend=True
+        )
+        plt.savefig(filename, format="png")
+        plt.close()
+
+        self.assertTrue(isfile(filename))
 
     def test_plot_network_options(self):
         # NOTE:to compare with the old plot_network set compare=True.
@@ -126,7 +172,7 @@ def test_plot_network_options(self):
         cmap = matplotlib.colormaps['viridis']
 
 
-        inp_file = join(ex_datadir, "Net3.inp")
+        inp_file = join(ex_datadir, "Net6.inp")
         wn = wntr.network.WaterNetworkModel(inp_file)
 
         random_node_values = pd.Series(
@@ -171,8 +217,8 @@ def test_plot_network_options(self):
              "link_cmap": cmap,
              "link_range": [0,1],
              "link_width": 1.5},
-            {"show_pump_direction": True,
-             "show_pump_direction": True}
+            # {"show_pump_direction": True,
+            #  "show_pump_direction": True}
         ]
 
         for kwargs in kwarg_list: