def test_worker_stop_timeout(self):
class MyTestUser(User):
_test_state = 0
wait_time = constant(0)
@task
def the_task(self):
MyTestUser._test_state = 1
gevent.sleep(0.2)
MyTestUser._test_state = 2
with mock.patch("locust.rpc.rpc.Client", mocked_rpc()) as client:
environment = Environment()
test_start_run = [False]
@environment.events.test_start.add_listener
def on_test_start(**kw):
test_start_run[0] = True
worker = self.get_runner(environment=environment, user_classes=[MyTestUser])
self.assertEqual(1, len(client.outbox))
self.assertEqual("client_ready", client.outbox[0].type)
client.mocked_send(Message("hatch", {
"hatch_rate": 1,
"num_users": 1,
"host": "",
"stop_timeout": 1,
}, "dummy_client_id"))
#print("outbox:", client.outbox)
# wait for worker to hatch locusts
self.assertIn("hatching", [m.type for m in client.outbox])
worker.hatching_greenlet.join()
self.assertEqual(1, len(worker.user_greenlets))
# check that locust has started running
gevent.sleep(0.01)
self.assertEqual(1, MyTestUser._test_state)
# send stop message
client.mocked_send(Message("stop", None, "dummy_client_id"))
worker.user_greenlets.join()
# check that locust user got to finish
self.assertEqual(2, MyTestUser._test_state)
# make sure the test_start was never fired on the worker
self.assertFalse(test_start_run[0])
def test_worker_stop_timeout(self):
class MyTestLocust(Locust):
_test_state = 0
class task_set(TaskSet):
wait_time = constant(0)
@task
def the_task(self):
MyTestLocust._test_state = 1
gevent.sleep(0.2)
MyTestLocust._test_state = 2
with mock.patch("locust.rpc.rpc.Client", mocked_rpc()) as client:
environment = Environment(options=mocked_options())
worker = self.get_runner(environment=environment,
locust_classes=[MyTestLocust])
self.assertEqual(1, len(client.outbox))
self.assertEqual("client_ready", client.outbox[0].type)
client.mocked_send(
Message(
"hatch", {
"hatch_rate": 1,
"num_clients": 1,
"host": "",
"stop_timeout": 1,
}, "dummy_client_id"))
#print("outbox:", client.outbox)
# wait for worker to hatch locusts
self.assertIn("hatching", [m.type for m in client.outbox])
worker.hatching_greenlet.join()
self.assertEqual(1, len(worker.locusts))
# check that locust has started running
gevent.sleep(0.01)
self.assertEqual(1, MyTestLocust._test_state)
# send stop message
client.mocked_send(Message("stop", None, "dummy_client_id"))
worker.locusts.join()
# check that locust user got to finish
self.assertEqual(2, MyTestLocust._test_state)
def test_slave_without_stop_timeout(self):
class MyTestLocust(Locust):
_test_state = 0
class task_set(TaskSet):
wait_time = constant(0)
@task
def the_task(self):
MyTestLocust._test_state = 1
gevent.sleep(0.2)
MyTestLocust._test_state = 2
with mock.patch("locust.rpc.rpc.Client", mocked_rpc()) as client:
options = mocked_options()
options.stop_timeout = None
slave = SlaveLocustRunner([MyTestLocust], options)
self.assertEqual(1, len(client.outbox))
self.assertEqual("client_ready", client.outbox[0].type)
client.mocked_send(
Message(
"hatch", {
"hatch_rate": 1,
"num_clients": 1,
"host": "",
"stop_timeout": None,
}, "dummy_client_id"))
#print("outbox:", client.outbox)
# wait for slave to hatch locusts
self.assertIn("hatching", [m.type for m in client.outbox])
slave.hatching_greenlet.join()
self.assertEqual(1, len(slave.locusts))
# check that locust has started running
gevent.sleep(0.01)
self.assertEqual(1, MyTestLocust._test_state)
# send stop message
client.mocked_send(Message("stop", None, "dummy_client_id"))
slave.locusts.join()
# check that locust user did not get to finish
self.assertEqual(1, MyTestLocust._test_state)
def test_slave_stats_report_median(self):
import mock
class MyTestLocust(Locust):
pass
with mock.patch("locust.rpc.rpc.Server",
mocked_rpc_server()) as server:
master = MasterLocustRunner(MyTestLocust, self.options)
server.mocked_send(Message("client_ready", None, "fake_client"))
master.stats.get("/", "GET").log(100, 23455)
master.stats.get("/", "GET").log(800, 23455)
master.stats.get("/", "GET").log(700, 23455)
data = {"user_count": 1}
events.report_to_master.fire(client_id="fake_client", data=data)
master.stats.clear_all()
server.mocked_send(Message("stats", data, "fake_client"))
s = master.stats.get("/", "GET")
self.assertEqual(700, s.median_response_time)
def heartbeat(self):
while True:
try:
self.client.send(
Message(
'heartbeat', {
'state': self.worker_state,
'current_cpu_usage': self.current_cpu_usage
}, self.client_id))
except RPCError as e:
logger.error("RPCError found when sending heartbeat: %s" % (e))
self.reset_connection()
gevent.sleep(HEARTBEAT_INTERVAL)
def test_stats_reporting(self):
import mock
class MyTestLocust(Locust):
pass
with mock.patch("locust.rpc.rpc.SlaveServer",
mocked_rpc_server()) as server:
with mock.patch("locust.rpc.rpc.SlaveClient",
mocked_rpc_server()) as client:
with mock.patch("locust.runners.slave.Process",
mocked_process()) as processes:
self.slave = SlaveLocustRunner(MyTestLocust,
config.locust_config())
server.mocked_send(
'all', Message("worker_ready", None, "fake_client0"))
self.slave.start_hatching(1, 1)
self.slave.stats.get("Task", "/", "GET").log(100, 23455)
self.slave.stats.get("Task", "/", "GET").log(800, 23455)
self.slave.stats.get("Task", "/", "GET").log(700, 23455)
sleep(runners.slave.SLAVE_STATS_INTERVAL + 0.1)
messages = [
Message.unserialize(msg) for msg in client.outbox_all
]
data = filter(
lambda m: m.type == 'stats' and m.data['stats'],
messages)[0].data
self.assertEqual({
800: 1,
100: 1,
700: 1
}, data['stats'][0]['response_times'])
self.assertEqual(3, data['stats'][0]['num_requests'])
self.assertEqual(1600,
data['stats'][0]['total_response_time'])
self.assertEqual(1, data['worker_count'])
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_slave_connect(self):
class MyTestLocust(Locust):
pass
with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
master = MasterLocustRunner(MyTestLocust, self.options)
server.mocked_send(
Message("client_ready", None, "zeh_fake_client1"))
self.assertEqual(1, len(master.clients))
self.assertTrue(
"zeh_fake_client1" in master.clients,
"Could not find fake client in master instance's clients dict")
server.mocked_send(
Message("client_ready", None, "zeh_fake_client2"))
server.mocked_send(
Message("client_ready", None, "zeh_fake_client3"))
server.mocked_send(
Message("client_ready", None, "zeh_fake_client4"))
self.assertEqual(4, len(master.clients))
server.mocked_send(Message("quit", None, "zeh_fake_client3"))
self.assertEqual(3, len(master.clients))
def worker(self):
while True:
msg = self.client.recv()
if msg.type == "hatch":
self.client.send(Message("hatching", None, self.client_id))
job = msg.data
self.hatch_rate = job["hatch_rate"]
self.num_requests = job["num_requests"]
self.host = job["host"]
self.hatching_greenlet = gevent.spawn(
lambda: self.start_hatching(locust_count=job["num_clients"
],
hatch_rate=job["hatch_rate"]))
elif msg.type == "stop":
self.stop()
self.client.send(
Message("client_stopped", None, self.client_id))
self.client.send(Message("client_ready", None, self.client_id))
elif msg.type == "quit":
logger.info("Got quit message from master, shutting down...")
self.stop()
self.greenlet.kill(block=True)
def test_master_total_stats(self):
with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
master = self.get_runner()
server.mocked_send(Message("client_ready", None, "fake_client"))
stats = RequestStats()
stats.log_request("GET", "/1", 100, 3546)
stats.log_request("GET", "/1", 800, 56743)
stats2 = RequestStats()
stats2.log_request("GET", "/2", 700, 2201)
server.mocked_send(Message("stats", {
"stats":stats.serialize_stats(),
"stats_total": stats.total.serialize(),
"errors":stats.serialize_errors(),
"user_count": 1,
}, "fake_client"))
server.mocked_send(Message("stats", {
"stats":stats2.serialize_stats(),
"stats_total": stats2.total.serialize(),
"errors":stats2.serialize_errors(),
"user_count": 2,
}, "fake_client"))
self.assertEqual(700, master.stats.total.median_response_time)
def test_change_user_count_during_hatching(self):
class User(Locust):
wait_time = constant(1)
class task_set(TaskSet):
@task
def my_task(self):
pass
with mock.patch("locust.rpc.rpc.Client", mocked_rpc()) as client:
options = mocked_options()
options.stop_timeout = None
environment = Environment(options=options)
slave = self.get_runner(environment=environment,
locust_classes=[User])
client.mocked_send(
Message(
"hatch", {
"hatch_rate": 5,
"num_clients": 10,
"host": "",
"stop_timeout": None,
}, "dummy_client_id"))
sleep(0.6)
self.assertEqual(STATE_HATCHING, slave.state)
client.mocked_send(
Message(
"hatch", {
"hatch_rate": 5,
"num_clients": 9,
"host": "",
"stop_timeout": None,
}, "dummy_client_id"))
sleep(0)
slave.hatching_greenlet.join()
self.assertEqual(9, len(slave.locusts))
slave.quit()
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
"""
import mock
class MyTestLocust(Locust):
pass
with mock.patch("locust.rpc.rpc.SlaveServer",
mocked_rpc_server()) as server:
with mock.patch("locust.runners.slave.Process", mocked_process()):
self.slave = SlaveLocustRunner(MyTestLocust,
config.locust_config())
server.mocked_send(
'all', Message("worker_ready", None, "fake_client0"))
sleep(0)
self.slave.start_hatching(1, 1)
for i in range(1, 5):
server.mocked_send(
"all",
Message("worker_ready", None, "fake_client%i" % i))
sleep(0)
del server.outbox_direct[:]
self.slave.start_hatching(42, 7)
self.assertEqual(5, len(server.outbox_direct))
num_clients = 0
for msg in server.outbox_direct:
num_clients += Message.unserialize(
msg[1]).data["num_clients"]
self.assertEqual(
42, num_clients,
"Total number of locusts that would have been spawned is not 42"
)
def test_rebalance_locust_users_on_worker_connect(self):
with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
master = self.get_runner()
server.mocked_send(Message("client_ready", None, "zeh_fake_client1"))
self.assertEqual(1, len(master.clients))
self.assertTrue("zeh_fake_client1" in master.clients, "Could not find fake client in master instance's clients dict")
master.start(100, 20)
self.assertEqual(1, len(server.outbox))
client_id, msg = server.outbox.pop()
self.assertEqual(100, msg.data["num_users"])
self.assertEqual(20, msg.data["hatch_rate"])
# let another worker connect
server.mocked_send(Message("client_ready", None, "zeh_fake_client2"))
self.assertEqual(2, len(master.clients))
self.assertEqual(2, len(server.outbox))
client_id, msg = server.outbox.pop()
self.assertEqual(50, msg.data["num_users"])
self.assertEqual(10, msg.data["hatch_rate"])
client_id, msg = server.outbox.pop()
self.assertEqual(50, msg.data["num_users"])
self.assertEqual(10, msg.data["hatch_rate"])
def test_spawn_fewer_locusts_than_workers(self):
with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
master = self.get_runner()
for i in range(5):
server.mocked_send(Message("client_ready", None, "fake_client%i" % i))
master.start(2, 2)
self.assertEqual(5, len(server.outbox))
num_users = 0
for _, msg in server.outbox:
num_users += msg.data["num_users"]
self.assertEqual(2, num_users, "Total number of locusts that would have been spawned is not 2")
def worker_relaunch(self):
import mock
class MyTestLocust(Locust):
pass
with mock.patch("locust.rpc.rpc.SlaveServer",
mocked_rpc_server()) as server:
with mock.patch("locust.rpc.rpc.SlaveClient",
mocked_rpc_server()) as client:
with mock.patch("locust.runners.slave.Process",
mocked_process()) as processes:
self.slave = SlaveLocustRunner(MyTestLocust,
config.locust_config())
server.mocked_send(
'all', Message("worker_ready", None, "fake_client0"))
self.slave.start_hatching(1, 1)
self.assertEqual(1, len(processes.started))
sleep(2 * runners.slave.HEARTBEAT_INTERVAL)
self.assertEqual(0, self.slave.worker_count)
client.mocked_send('all', Message("ping", None, "master"))
sleep(runners.slave.HEARTBEAT_INTERVAL)
self.assertEqual(2, len(processes.started))
def worker(self):
while True:
try:
msg = self.client.recv()
except RPCError as e:
logger.error("RPCError found when receiving from master: %s" %
(e))
continue
if msg.type == "spawn":
self.worker_state = STATE_SPAWNING
self.client.send(Message("spawning", None, self.client_id))
job = msg.data
self.spawn_rate = job["spawn_rate"]
self.target_user_count = job["num_users"]
self.environment.host = job["host"]
self.environment.stop_timeout = job["stop_timeout"]
if self.spawning_greenlet:
# kill existing spawning greenlet before we launch new one
self.spawning_greenlet.kill(block=True)
self.spawning_greenlet = self.greenlet.spawn(
lambda: self.start(user_count=job["num_users"],
spawn_rate=job["spawn_rate"]))
self.spawning_greenlet.link_exception(
greenlet_exception_handler)
elif msg.type == "stop":
self.stop()
self.client.send(
Message("client_stopped", None, self.client_id))
self.client.send(Message("client_ready", None, self.client_id))
self.worker_state = STATE_INIT
elif msg.type == "quit":
logger.info("Got quit message from master, shutting down...")
self.stop()
self._send_stats(
) # send a final report, in case there were any samples not yet reported
self.greenlet.kill(block=True)
def test_on_ping(self):
import mock
class MyTestLocust(Locust):
pass
with mock.patch("locust.rpc.rpc.WorkerClient",
mocked_rpc_server()) as client:
self.worker = WorkerLocustRunner(MyTestLocust,
config.locust_config())
sleep(0)
self.assertEqual(2, len(client.outbox_all))
client.mocked_send('all', Message("ping", None, "master"))
sleep(0)
self.assertEqual(3, len(client.outbox_all))
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
"""
import mock
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))
sleep(0)
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 __init__(self, locust_classes, options):
super(WorkerLocustRunner, self).__init__(locust_classes, options)
self.master_port = options.master_port
self.worker_id = socket.gethostname().replace(
' ', '_') + "_" + uuid.uuid4().hex
self.greenlet = Group()
self.client = rpc.WorkerClient(self.master_port, self.worker_id)
self.client.bind_handler(self.WorkerClientHandler(self))
def noop(*args, **kwargs):
pass
# register listener for when all locust users have hatched, and report it to the master node
def on_hatch_complete(user_count):
self.client.send_all(
Message("hatch_complete", {"count": user_count},
self.worker_id))
events.hatch_complete += on_hatch_complete
# register listener that adds the current number of spawned locusts to
# the report that is sent to the master node
def on_report_to_master(node_id, data):
data["user_count"] = self.user_count
events.report_to_master += on_report_to_master
# register listener that sends quit message to master
def on_quitting():
self.quit()
events.quitting += on_quitting
# register listener thats sends locust exceptions to master
def on_locust_error(locust_instance, exception, tb):
formatted_tb = "".join(traceback.format_tb(tb))
data = {"msg": str(exception), "traceback": formatted_tb}
self.client.send_all(Message("exception", data, self.worker_id))
events.locust_error += on_locust_error
self.greenlet.spawn(self.slave_listener).link_exception(callback=noop)
gevent.sleep(0.5)
self.client.send_all(Message("worker_ready", None, self.worker_id))
self.greenlet.spawn(self.stats_reporter).link_exception(callback=noop)
def test_on_ping(self):
import mock
class MyTestLocust(Locust):
pass
with mock.patch("locust.rpc.rpc.SlaveClient",
mocked_rpc_server()) as client:
with mock.patch("locust.runners.slave.Process",
mocked_process()) as processes:
self.slave = SlaveLocustRunner(MyTestLocust,
config.locust_config())
sleep(0)
self.assertEqual(2, len(client.outbox_all))
client.mocked_send('all', Message("ping", None, "master"))
sleep(0)
self.assertEqual(3, len(client.outbox_all))
def __init__(self, *args, **kwargs):
super(SlaveLocustRunner, self).__init__(*args, **kwargs)
random_hash = md5(
str(time() +
random.randint(0, 10000)).encode('utf-8')).hexdigest()
self.client_id = socket.gethostname() + "_" + random_hash
self.client = rpc.Client(self.master_host, self.master_port)
self.greenlet = Group()
self.greenlet.spawn(self.worker).link_exception(callback=self.noop)
self.client.send(Message("client_ready", None, self.client_id))
self.greenlet.spawn(
self.stats_reporter).link_exception(callback=self.noop)
# Register listener for when all locust users have hatched, and report it to the master node
def on_hatch_complete(user_count):
self.client.send(
Message("hatch_complete", {"count": user_count},
self.client_id))
events.hatch_complete += on_hatch_complete
# Register listener that adds the current number of spawned locusts to the report
# that is sent to the master node
def on_report_to_master(_, data):
data["user_count"] = self.user_count
events.report_to_master += on_report_to_master
# register listener that sends quit message to master
def on_quitting():
self.client.send(Message("quit", None, self.client_id))
events.quitting += on_quitting
# register listener thats sends locust exceptions to master
def on_locust_error(locust_instance, exception, tb):
formatted_tb = "".join(traceback.format_tb(tb))
self.client.send(
Message("exception", {
"msg": str(exception),
"traceback": formatted_tb
}, self.client_id))
events.locust_error += on_locust_error
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 start(self, user_count, spawn_rate):
self.target_user_count = user_count
num_workers = len(self.clients.ready) + len(
self.clients.running) + len(self.clients.spawning)
if not num_workers:
logger.warning(
"You are running in distributed mode but have no worker servers connected. "
"Please connect workers prior to swarming.")
return
self.spawn_rate = spawn_rate
worker_num_users = user_count // (num_workers or 1)
worker_spawn_rate = float(spawn_rate) / (num_workers or 1)
remaining = user_count % num_workers
logger.info(
"Sending spawn jobs of %d users and %.2f spawn rate to %d ready clients"
% (worker_num_users, worker_spawn_rate, num_workers))
if worker_spawn_rate > 100:
logger.warning(
"Your selected spawn rate is very high (>100/worker), and this is known to sometimes cause issues. Do you really need to ramp up that fast?"
)
if self.state != STATE_RUNNING and self.state != STATE_SPAWNING:
self.stats.clear_all()
self.exceptions = {}
self.environment.events.test_start.fire(
environment=self.environment)
for client in (self.clients.ready + self.clients.running +
self.clients.spawning):
data = {
"spawn_rate": worker_spawn_rate,
"num_users": worker_num_users,
"host": self.environment.host,
"stop_timeout": self.environment.stop_timeout,
}
if remaining > 0:
data["num_users"] += 1
remaining -= 1
self.server.send_to_client(Message("spawn", data, client.id))
self.state = STATE_SPAWNING
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 workers
"""
with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
master = self.get_runner()
for i in range(5):
server.mocked_send(Message("client_ready", None, "fake_client%i" % i))
master.start(7, 7)
self.assertEqual(5, len(server.outbox))
num_users = 0
for _, msg in server.outbox:
num_users += msg.data["num_users"]
self.assertEqual(7, num_users, "Total number of locusts that would have been spawned is not 7")
def test_stop_event_quit(self):
"""
Tests that test_stop event is fired when quit() is called directly
"""
with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
master = self.get_runner()
run_count = [0]
@self.environment.events.test_stop.add_listener
def on_test_stop(*a, **kw):
run_count[0] += 1
for i in range(5):
server.mocked_send(Message("client_ready", None, "fake_client%i" % i))
master.start(7, 7)
self.assertEqual(5, len(server.outbox))
master.quit()
self.assertEqual(1, run_count[0])
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 test_spawn_locusts_in_stepload_mode(self):
class MyTestLocust(Locust):
pass
with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
master = MasterLocustRunner(MyTestLocust, self.options)
for i in range(5):
server.mocked_send(
Message("client_ready", None, "fake_client%i" % i))
# start a new swarming in Step Load mode: total locust count of 10, hatch rate of 2, step locust count of 5, step duration of 5s
master.start_stepload(10, 2, 5, 5)
# make sure the first step run is started
sleep(1)
self.assertEqual(5, len(server.outbox))
num_clients = 0
end_of_last_step = len(server.outbox)
for _, msg in server.outbox:
num_clients += msg.data["num_clients"]
self.assertEqual(
5, num_clients,
"Total number of locusts that would have been spawned for first step is not 5"
)
# make sure the first step run is complete
sleep(5)
num_clients = 0
idx = end_of_last_step
while idx < len(server.outbox):
msg = server.outbox[idx][1]
num_clients += msg.data["num_clients"]
idx += 1
self.assertEqual(
10, num_clients,
"Total number of locusts that would have been spawned for second step is not 10"
)
def test_start_event(self):
"""
Tests that test_start event is fired
"""
with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
master = self.get_runner()
run_count = [0]
@self.environment.events.test_start.add_listener
def on_test_start(*a, **kw):
run_count[0] += 1
for i in range(5):
server.mocked_send(Message("client_ready", None, "fake_client%i" % i))
master.start(7, 7)
self.assertEqual(5, len(server.outbox))
self.assertEqual(1, run_count[0])
# change number of users and check that test_start isn't fired again
master.start(7, 7)
self.assertEqual(1, run_count[0])
def test_worker_receive_propagated_config(self):
import mock
class MyTestLocust(Locust):
pass
with mock.patch("locust.rpc.rpc.SlaveClient",
mocked_rpc_server()) as client:
with mock.patch("locust.runners.slave.Process",
mocked_process()) as processes:
self.slave = SlaveLocustRunner(MyTestLocust,
config.locust_config())
sleep(0)
self.assertEqual(2, len(client.outbox_all))
msg = {
'host': 'custom_host.com',
'master_host': 'new_master_host.com'
}
client.mocked_send('all', Message("new_config", msg, "master"))
sleep(0)
self.assertEqual(self.slave.options.host,
'http://custom_host.com')
self.assertEqual(self.slave.options.master_host, '127.0.0.1')
def start_hatching(self, locust_count, hatch_rate):
worker_num = self.slave_count
if self.state != STATE.INIT and self.state != STATE.STOPPED:
self.stop()
self.stats.clear_all()
self.stats.reset_all()
self.exceptions = {}
events.master_start_hatching.fire()
self.state = STATE.HATCHING
leftover = locust_count - (locust_count / worker_num) * worker_num
adjust = lambda x: 1 if x <= leftover else 0
calc = lambda x: locust_count / worker_num + adjust(x)
slave_locust_count = [calc(x) for x in range(1, worker_num + 1)]
slave_hatch_rate = hatch_rate / float(worker_num)
if self.options.num_requests:
slave_num_requests = int(self.options.num_requests /
self.slave_count)
else:
slave_num_requests = None
logger.info("Sending hatch jobs to %d ready slave(s)", worker_num)
for client_rate, slave_id in zip(slave_locust_count,
self.slaves.keys()):
data = {
"hatch_rate": slave_hatch_rate,
"num_clients": client_rate,
"num_requests": slave_num_requests
}
self.server.send_to(slave_id, Message("hatch", data, None))
self.slaves[slave_id].task = data
self.stats.start_time = time.time()
self.state = STATE.HATCHING
def start_hatching(self, locust_count=None, hatch_rate=None, wait=False):
num_slaves = len(self.clients.ready) + len(self.clients.running)
if not num_slaves:
logger.warning(
"You are running in distributed mode but have no slave servers connected. "
"Please connect slaves prior to swarming.")
return
self.num_clients = locust_count
slave_num_clients = locust_count // (num_slaves or 1)
slave_hatch_rate = float(hatch_rate) / (num_slaves or 1)
remaining = locust_count % num_slaves
logger.info("Sending hatch jobs to %d ready clients", num_slaves)
if self.state != STATE_RUNNING and self.state != STATE_HATCHING:
self.stats.clear_all()
self.exceptions = {}
events.master_start_hatching.fire()
for _ in six.itervalues(self.clients):
data = {
"hatch_rate": slave_hatch_rate,
"num_clients": slave_num_clients,
"num_requests": self.num_requests,
"host": self.host,
"stop_timeout": None
}
if remaining > 0:
data["num_clients"] += 1
remaining -= 1
self.server.send(Message("hatch", data, None))
self.stats.start_time = time()
self.state = STATE_HATCHING
def test_last_worker_quitting_stops_test(self):
with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
master = self.get_runner()
server.mocked_send(Message("client_ready", None, "fake_client1"))
server.mocked_send(Message("client_ready", None, "fake_client2"))
master.start(1, 2)
server.mocked_send(Message("hatching", None, "fake_client1"))
server.mocked_send(Message("hatching", None, "fake_client2"))
server.mocked_send(Message("quit", None, "fake_client1"))
sleep(0)
self.assertEqual(1, len(master.clients.all))
self.assertNotEqual(STATE_STOPPED, master.state, "Not all workers quit but test stopped anyway.")
server.mocked_send(Message("quit", None, "fake_client2"))
sleep(0)
self.assertEqual(0, len(master.clients.all))
self.assertEqual(STATE_STOPPED, master.state, "All workers quit but test didn't stop.")
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
