From ec3dff12b0a85955f07b29d5c6dbe0fe424f376e Mon Sep 17 00:00:00 2001
From: wqy <wqyuyu@outlook.com>
Date: Sat, 17 Sep 2022 22:57:42 +0800
Subject: [PATCH] Add pump test case

---
 api/s6_pumps.py   |  2 +-
 test_tjnetwork.py | 81 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+), 1 deletion(-)

diff --git a/api/s6_pumps.py b/api/s6_pumps.py
index 7fa40ab..21d9569 100644
--- a/api/s6_pumps.py
+++ b/api/s6_pumps.py
@@ -12,7 +12,7 @@ def add_pump(name: str, id: str, node1: str, node2: str) -> ChangeSet:
     if node1 == node2:
         return ChangeSet()
 
-    sql = f"insert into pumps (id, node1, node2) values ('{id}', '{node1}', '{node2}'');"
+    sql = f"insert into pumps (id, node1, node2) values ('{id}', '{node1}', '{node2}');"
     undo_sql = f'delete from pumps where id = "{id}";'
     return add_link(name, PUMP, id, sql, undo_sql)
 
diff --git a/test_tjnetwork.py b/test_tjnetwork.py
index 73d2629..6bdf309 100644
--- a/test_tjnetwork.py
+++ b/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):