Fix tests after changing ZMQ sockets to DEALER-BROKER

locust/test/test_runners.py

@@ -35,6 +35,14 @@ def recv(self):
         def send(self, message):
             self.outbox.append(message.serialize())
 
+        def send_to_client(self, message):
+            self.outbox.append([message.node_id, message.serialize()])
+
+        def recv_from_client(self):
+            results = self.queue.get()
+            msg = Message.unserialize(results)
+            return msg.node_id, msg
+
     return MockedRpcServer
 
 
@@ -218,7 +226,7 @@ class MyTestLocust(Locust):
 
             num_clients = 0
             for msg in server.outbox:
-                num_clients += Message.unserialize(msg).data["num_clients"]
+                num_clients += Message.unserialize(msg[1]).data["num_clients"]
 
             self.assertEqual(7, num_clients, "Total number of locusts that would have been spawned is not 7")
 
@@ -238,7 +246,7 @@ class MyTestLocust(Locust):
 
             num_clients = 0
             for msg in server.outbox:
-                num_clients += Message.unserialize(msg).data["num_clients"]
+                num_clients += Message.unserialize(msg[1]).data["num_clients"]
 
             self.assertEqual(2, num_clients, "Total number of locusts that would have been spawned is not 2")