Add pump test case

test_tjnetwork.py

@@ -368,6 +368,87 @@ class TestApi:
         self.leave(p)
 
     def test_pump(self):
+        p = "test_pump"
+        self.enter(p)
+
+        change_set = add_pump(p, 'p1', 'j1', 'j2')
+        assert len(change_set.added) == 0
+
+        add_junction(p, 'j1', 0.0, 0.0, 0.0)
+        add_junction(p, 'j2', 10.0, 10.0, 0.0)
+        add_junction(p, 'j3', 0.0, 0.0, 0.0)
+        add_junction(p, 'j4', 10.0, 10.0, 0.0)
+        assert is_junction(p, 'j1')
+        assert is_junction(p, 'j2')
+        assert is_junction(p, 'j3')
+        assert is_junction(p, 'j4')
+
+        assert not is_link(p, 'p1')
+
+        change_set = add_pump(p, 'p1', 'j1', 'j1')
+        assert len(change_set.added) == 0
+
+        change_set = add_pump(p, 'p1', 'j1', 'j2')
+        assert len(change_set.added) == 1
+        assert change_set.added[0]['type'] == PUMP
+        assert change_set.added[0]['id'] == "p1"
+
+        change_set = add_pump(p, 'p1', 'j1', 'j2')
+        assert len(change_set.added) == 0
+
+        pumps = get_links(p)
+        assert len(pumps) == 1
+
+        assert get_pump_node1(p, 'p1') == 'j1'
+        assert get_pump_node2(p, 'p1') == 'j2'
+
+        change_set = set_pump_node1(p, 'p1', 'j2')
+        assert len(change_set.updated) == 0
+        change_set = set_pump_node2(p, 'p1', 'j1')
+        assert len(change_set.updated) == 0
+
+        change_set = set_pump_node1(p, 'p1', 'j3')
+        assert len(change_set.updated) == 1
+        assert change_set.updated[0]['type'] == PUMP
+        assert change_set.updated[0]['id'] == 'p1'
+        assert change_set.updated[0]['property'] == 'node1'
+
+        change_set = set_pump_node2(p, 'p1', 'j4')
+        assert len(change_set.updated) == 1
+        assert change_set.updated[0]['type'] == PUMP
+        assert change_set.updated[0]['id'] == 'p1'
+        assert change_set.updated[0]['property'] == 'node2'
+
+        assert get_pump_node1(p, 'p1') == 'j3'
+        assert get_pump_node2(p, 'p1') == 'j4'
+
+        change_set = add_pump(p, 'p2', 'j1', 'j2')
+        assert len(change_set.added) == 1
+        assert change_set.added[0]['type'] == PUMP
+        assert change_set.added[0]['id'] == "p2"
+
+        pumps = get_links(p)
+        assert len(pumps) == 2
+
+        change_set = delete_pump(p, 'p')
+        assert len(change_set.deleted) == 0
+
+        assert is_pump(p, 'p1')
+
+        change_set = delete_pump(p, 'p1')
+        assert len(change_set.deleted) == 1
+        assert change_set.deleted[0]['type'] == PUMP
+        assert change_set.deleted[0]['id'] == "p1"
+
+        pumps = get_links(p)
+        assert len(pumps) == 1
+
+        change_set = delete_pump(p, 'p2')
+        assert len(change_set.deleted) == 1
+        assert change_set.deleted[0]['type'] == PUMP
+        assert change_set.deleted[0]['id'] == "p2"
+
+        self.leave(p)
         pass
 
     def test_valve(self):