将 Locust 作为库使用

可以通过自己的 Python 代码启动负载测试,而不是使用 locust 命令运行 Locust。

首先,创建一个 Environment 实例:

from locust.env import Environment

env = Environment(user_classes=[MyTestUser])

然后,可以使用 Environment 实例的 create_local_runnercreate_master_runner 方法启动一个 Runner 实例,用于启动负载测试:

env.create_local_runner()
env.runner.start(5000, spawn_rate=20)
env.runner.greenlet.join()

也可以绕过调度和分发逻辑,手动控制启动的用户:

new_users = env.runner.spawn_users({MyUserClass.__name__: 2})
new_users[1].my_custom_token = "custom-token-2"
new_users[0].my_custom_token = "custom-token-1"

上述示例仅在独立/本地运行模式下有效,并且是一个实验性特性。一个更常见/更好的方法是使用 inittest_start 事件钩子来获取/创建令牌列表,并在 on_starton_stop 方法中从该列表中读取并设置它们到各个用户实例。

尽管可以通过 create_worker_runner 创建 locust 工作者(worker),但几乎没有实际意义。每个工作者需要在单独的 Python 进程中运行,直接与工作者运行器交互可能会导致问题。最好使用常规的 locust --worker …​ 命令启动工作者。

还可以使用 Environment 实例的 create_web_ui 方法启动一个 Web UI,用于查看统计信息,并控制运行器(例如启动和停止负载测试):

env.create_local_runner()
env.create_web_ui()
env.web_ui.greenlet.join()

跳过猴子补丁

一些包(例如 boto3)在使用 Locust 作为库时可能存在兼容性问题,特别是当猴子补丁已经应用时。在这种情况下,可以通过设置环境变量 LOCUST_SKIP_MONKEY_PATCH=1 来禁用猴子补丁。

完整示例

#!/usr/bin/env python3
from locust import HttpUser, events, task
from locust.env import Environment
from locust.log import setup_logging
from locust.stats import stats_history, stats_printer

import gevent

setup_logging("INFO")


class MyUser(HttpUser):
    host = "https://docs.locust.io"

    @task
    def t(self):
        self.client.get("/")


# 设置 Environment 和 Runner
env = Environment(user_classes=[MyUser], events=events)
runner = env.create_local_runner()

# 启动 WebUI 实例
web_ui = env.create_web_ui("127.0.0.1", 8089)

# 执行初始化事件处理程序(如果已注册)
env.events.init.fire(environment=env, runner=runner, web_ui=web_ui)

# 启动一个 greenlet 定期输出当前统计信息
gevent.spawn(stats_printer, env.stats)

# 启动一个 greenlet 保存当前统计信息到历史记录
gevent.spawn(stats_history, env.runner)

# 启动测试
runner.start(1, spawn_rate=10)

# 30 秒后停止运行器
gevent.spawn_later(30, runner.quit)

# 等待 greenlet 结束
runner.greenlet.join()

# 停止 Web 服务器以防万一
web_ui.stop()