def test_spawn_fewer_locusts_than_slaves(self):
class MyTestLocust(Locust):
pass
with mock.patch("locust.rpc.rpc.Server", mocked_rpc_server()) as server:
master = MasterLocustRunner(MyTestLocust, self.options)
for i in range(5):
server.mocked_send(Message("client_ready", None, "fake_client%i" % i))
master.start_hatching(2, 2)
self.assertEqual(5, len(server.outbox))
num_clients = 0
for _, msg in server.outbox:
num_clients += Message.unserialize(msg).data["num_clients"]
self.assertEqual(2, num_clients, "Total number of locusts that would have been spawned is not 2")
def test_spawn_uneven_locusts(self):
"""
Tests that we can accurately spawn a certain number of locusts, even if it's not an
even number of the connected slaves
"""
class MyTestLocust(Locust):
pass
with mock.patch("locust.rpc.rpc.Server", mocked_rpc_server()) as server:
master = MasterLocustRunner(MyTestLocust, self.options)
for i in range(5):
server.mocked_send(Message("client_ready", None, "fake_client%i" % i))
master.start_hatching(7, 7)
self.assertEqual(5, len(server.outbox))
num_clients = 0
for _, msg in server.outbox:
num_clients += Message.unserialize(msg).data["num_clients"]
self.assertEqual(7, num_clients, "Total number of locusts that would have been spawned is not 7")
def recv(self):
results = self.queue.get()
return Message.unserialize(results)
def test_message_serialize(self):
msg = Message("client_ready", None, "my_id")
rebuilt = Message.unserialize(msg.serialize())
self.assertEqual(msg.type, rebuilt.type)
self.assertEqual(msg.data, rebuilt.data)
self.assertEqual(msg.node_id, rebuilt.node_id)
def recv_from_client(self):
results = self.queue.get()
msg = Message.unserialize(results)
return msg.node_id, msg
def recv(self):
results = self.queue.get()
return Message.unserialize(results)
def test_message_serialize(self):
msg = Message("client_ready", None, "my_id")
rebuilt = Message.unserialize(msg.serialize())
self.assertEqual(msg.type, rebuilt.type)
self.assertEqual(msg.data, rebuilt.data)
self.assertEqual(msg.node_id, rebuilt.node_id)
def recv_from_client(self):
results = self.queue.get()
msg = Message.unserialize(results)
if msg.data == NETWORK_BROKEN:
raise RPCError()
return msg.node_id, msg
def recv_from_client(self):
results = self.queue.get()
msg = Message.unserialize(results)
return msg.node_id, msg
def recv(self):
data = self.queue.get()
msg = Message.unserialize(data.split(' ', 1)[1])
getattr(self.handler, 'on_' + msg.type)(msg)
