update tests

locustio · Mar 16, 2020 · f523314 · f523314
1 parent 8b26f51
commit f523314
Showing 1 changed file with 14 additions and 4 deletions.
diff --git a/locust/test/test_zmqrpc.py b/locust/test/test_zmqrpc.py
@@ -3,7 +3,7 @@
 import zmq
 from locust.rpc import zmqrpc, Message
 from locust.test.testcases import LocustTestCase
-
+from locust.exception import RPCError
 
 class ZMQRPC_tests(LocustTestCase):
     def setUp(self):
@@ -12,8 +12,8 @@ def setUp(self):
         self.client = zmqrpc.Client('localhost', self.server.port, 'identity')
 
     def tearDown(self):
-        self.server.socket.close()
-        self.client.socket.close()
+        self.server.close()
+        self.client.close()
         super(ZMQRPC_tests, self).tearDown()
 
     def test_constructor(self):
@@ -42,5 +42,15 @@ def test_client_recv(self):
     def test_client_retry(self):
         server = zmqrpc.Server('127.0.0.1', 0)
         server.socket.close()
-        with self.assertRaises(zmq.error.ZMQError):
+        with self.assertRaises(RPCError):
             server.recv_from_client()
+
+    def test_rpc_error(self):
+        server = zmqrpc.Server('127.0.0.1', 5557)
+        with self.assertRaises(RPCError):
+            server = zmqrpc.Server('127.0.0.1', 5557)
+        server.close()
+        with self.assertRaises(RPCError):
+            server.send_to_client(Message('test', 'message', 'identity'))
+
+