【Locust2.x】解决分布式测试时测试数据分发问题，跨节点通信，给每个worker分发不同的测试数据

1、初始化监听方法

• 如果是worker监听test_users、worker_info方法
• 如果是master监听acknowledge_users方法

@events.init.add_listener
def on_locust_init_testdata(environment, **_kwargs):
    if not isinstance(environment.runner, MasterRunner):
        # 初始化locust，注册test_users方法，可以接收到test_users的消息，并且把消息内容给setup_test_users处理
        environment.runner.register_message("test_users", setup_test_users)
        environment.runner.register_message("worker_info", setup_worker_info)
    if not isinstance(environment.runner, WorkerRunner):
        environment.runner.register_message("acknowledge_users", on_acknowledge)

2、测试开始时分发测试数据

@events.test_start.add_listener
def on_test_start(environment, **_kwargs):
	# 测试开始前，如果是worker则不执行分发数据操作
    if isinstance(environment.runner, WorkerRunner):
        return
	# res是数据源，从文件中读取，可自行定义
    res = read_text_line("pro", "pro_stu")
    users = [{"name": res[i]} for i in range(len(res))]
    # 获取worker数量
    worker_count = environment.runner.worker_count
    # 根据worker数量计算每个worker需要的数据大小
    chunk_size = int(len(users) / worker_count)

    # environment.runner.clients 是一个列表，里面放的是每个worker的ID
    for i, worker in enumerate(environment.runner.clients):
        # 通过数据大小计算列表的开始、结束下标
        start_index = i * chunk_size
        end_index = start_index + chunk_size if i + 1 < worker_count else len(users)
        # 根据数据拿到的下标，截取数据
        data = users[start_index:end_index]
        # 发送消息给test_users，并且指定worker
        index_info = {"worker": worker, "用户数量": len(data), "start": start_index, "end": end_index}
        # 发送测试数据
        environment.runner.send_message("test_users", data, worker)
        # 发送详情信息
        environment.runner.send_message("worker_info", index_info, worker)

3、worker接收信息

# 在全局定义了一个队列，方便使用
usernames_queue = queue.Queue()

def setup_test_users(environment, msg, **kwargs):
    """ Worker 接收消息的方法, 在 on_locust_init 中定义 """
    # 接收发送的用户信息, 填充到队列中
    usernames.extend(map(lambda u: u["name"], msg.data))
    # 我这里是把消息放入了队列，如果使用可以自行定义
    for i in usernames:
        usernames_queue.put({"username": i})


def setup_worker_info(environment, msg, **kwargs):
    # msg是on_test_start发送过来的消息，把拿到的消息再发送给master
    environment.runner.send_message("acknowledge_users", f"worker info: {msg.data}")

4、master接收消息

def on_acknowledge(environment, msg, **kwargs):
    """ Master 接收消息的方法, 在 on_locust_init 中定义 """
    # 打印worker信息
    print(msg.data)

5、脚本中使用

class MyUser(HttpUser):
    @task
    def login(self):
        username = usernames_queue.get()["username"]
        payload = f'name={username}&passwd=XXXX'
        headers = {'Content-Type': 'application/x-www-form-urlencoded',}
        with self.client.post("/login", params=payload, headers=headers) as login_res:
        	token = login_res.json()["token"]
            ...
        usernames_queue.put({"username": username})