事件钩子

Locust 提供了多个事件钩子，可以用来以不同的方式扩展 Locust。

例如，以下是如何设置一个事件监听器，在请求完成后触发：

from locust import events

@events.request.add_listener
def my_request_handler(request_type, name, response_time, response_length, response,
                       context, exception, start_time, url, **kwargs):
    if exception:
        print(f"请求 {name} 失败，异常为 {exception}")
    else:
        print(f"成功发起请求到: {name}")
        print(f"响应内容为 {response.text}")

在上述示例中，通配符关键字参数 (**kwargs) 将为空，因为我们处理了所有参数，但它可以防止代码在 Locust 将来版本中添加新参数时出现问题。

此外，完全可以实现一个不提供此事件所有参数的客户端。例如，非 HTTP 协议可能没有 URL 或响应对象的概念。可以从监听函数定义中移除任何缺失的字段，或者使用默认参数。

在分布式模式下运行 Locust 时，在执行测试之前，可能需要在工作节点上进行一些设置。你可以通过检查节点的运行器类型来确认是否不在主节点上：

from locust import events
from locust.runners import MasterRunner

@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    if not isinstance(environment.runner, MasterRunner):
        print("开始进行测试设置")
    else:
        print("从主节点启动测试")

@events.test_stop.add_listener
def on_test_stop(environment, **kwargs):
    if not isinstance(environment.runner, MasterRunner):
        print("清理测试数据")
    else:
        print("从主节点停止测试")

你还可以使用事件来添加自定义命令行参数。

要查看可用事件的完整列表，请参阅事件钩子文档。

请求上下文

请求事件有一个 context 参数，可以让你传递关于请求的数据（如用户名、标签等）。它可以直接在请求方法的调用中设置，或者通过重写 User.context() 方法在用户级别设置。

通过请求方法设置上下文：

class MyUser(HttpUser):
    @task
    def t(self):
        self.client.post("/login", json={"username": "foo"})
        self.client.get("/other_request", context={"username": "foo"})

    @events.request.add_listener
    def on_request(context, **kwargs):
        if context:
            print(context["username"])

通过用户实例设置上下文：

class MyUser(HttpUser):
    def context(self):
        return {"username": self.username}

    @task
    def t(self):
        self.username = "foo"
        self.client.post("/login", json={"username": self.username})

    @events.request.add_listener
    def on_request(context, **kwargs):
        print(context["username"])

通过响应中的值设置上下文，使用 catch_response：

with self.client.get("/", catch_response=True) as resp:
    resp.request_meta["context"]["requestId"] = resp.json()["requestId"]

请求上下文不会改变 Locust 常规统计数据的计算方式。像 locust.cloud 这样的日志/报告解决方案使用上述机制将上下文保存到数据库中。

添加 Web 路由

Locust 使用 Flask 来提供 Web UI，因此可以轻松地向 Web UI 添加 Web 端点。通过监听 init 事件，我们可以获取 Flask 应用实例的引用，并使用它来设置一个新路由：

from locust import events

@events.init.add_listener
def on_locust_init(environment, **kw):
    @environment.web_ui.app.route("/added_page")
    def my_added_page():
        return "另一个页面"

现在你应该能够启动 Locust，并访问 http://127.0.0.1:8089/added_page 。注意，它不会自动作为新标签页添加，你需要直接输入该 URL。

扩展 Web UI

作为向 Web UI 添加简单 Web 路由的替代方法，你可以使用 Flask 蓝图和模板，不仅可以添加路由，还可以扩展 Web UI，允许你在内置 Locust 统计信息旁边显示自定义数据。这是一个更高级的用法，但可以大大增强 Web UI 的实用性和可定制性。

扩展 Web UI 的工作示例可以在 Locust 源代码的示例目录中找到。

extend_modern_web_ui.py：显示每个调用的 content-length 表格。
web_ui_cache_stats.py：显示每个调用的 Varnish Hit/Miss 统计信息。这个示例可以轻松扩展到其他 CDN 或缓存代理，并收集其他缓存统计信息，如缓存年龄、控制信息等。

为 Web UI 添加认证

当使用 --web-login 标志时，Locust 使用 Flask-Login 来处理认证。login_manager 会暴露在 environment.web_ui.app 中，这使得你可以实现任何你想要的认证方式！

要使用用户名/密码认证，只需为 environment.web_ui.auth_args 提供 username_password_callback。你需要负责定义回调的路由，并实现认证逻辑。

认证提供者还可以配置为允许通过 GitHub 或 SSO 提供商进行认证。只需提供所需认证提供者的列表。你可以指定标签和图标以显示在按钮上。callback_url 是按钮指向的 URL，你将负责定义回调路由以及与第三方的认证。

无论你是使用用户名/密码认证、认证提供者，还是两者结合，都需要提供 user_loader 给 login_manager。user_loader 应该在认证成功时返回 User 对象，在认证失败时返回 None。

要在登录页面上显示错误消息（例如用户名/密码组合错误），你可以将 auth_error 存储在会话对象中：session["auth_error"] = "用户名或密码错误"。如果你有无错误的消息要显示给用户，可以选择将 auth_info 存储在会话对象中：session["auth_info"] = "成功创建新用户！"。

有关完整示例，请参见认证示例。

在某些情况下，你可能希望进一步扩展认证表单中的字段。为此，向 environment.web_ui.auth_args 传递 custom_form 字典。在这种情况下，字段将由一系列输入框表示，callback_url 将由 custom_form.callback_url 配置，提交按钮也可以使用 custom_form.submit_button_text 进行配置。认证表单中的字段可以是文本框、选择框、复选框或密码字段。你还可以为特定字段验证重写 HTML 输入类型（例如：type=email）。

有关如何配置 custom_form 的完整示例，请参见认证示例。

运行后台 greenlet

由于 Locust 文件 “只是代码”，没有任何东西阻止你启动自己的 greenlet，以与实际的负载/用户并行运行。

例如，你可以监控测试的失败比例，并在超过某个阈值时停止运行：

import gevent
from locust import events
from locust.runners import STATE_STOPPING, STATE_STOPPED, STATE_CLEANUP, MasterRunner, LocalRunner

def checker(environment):
    while not environment.runner.state in [STATE_STOPPING, STATE_STOPPED, STATE_CLEANUP]:
        time.sleep(1)
        if environment.runner.stats.total.fail_ratio > 0.2:
            print(f"失败比例为 {environment.runner.stats.total.fail_ratio}, 正在退出")
            environment.runner.quit()
            return

@events.init.add_listener
def on_locust_init(environment, **_kwargs):
    # 仅在主节点或本地节点运行时启动检查器
    if isinstance(environment.runner, MasterRunner) or isinstance(environment.runner, LocalRunner):
        gevent.spawn(checker, environment)

Locust 文件参数化

有两种主要方式来参数化你的 Locust 文件。

基本的环境变量

像其他程序一样，你可以使用环境变量：

在 Linux/Mac 上：

MY_FUNKY_VAR=42 locust ...

在 Windows 上：

SET MY_FUNKY_VAR=42
locust ...

然后在 Locust 文件中访问它们：

import os
print(os.environ['MY_FUNKY_VAR'])

自定义命令行参数

你可以使用 init_command_line_parser 事件向 Locust 添加自定义命令行参数。自定义参数也会在 Web UI 中呈现并可编辑。如果指定了选项，它们将在 Web UI 中以下拉框的形式展示。

from locust import HttpUser, events, task

@events.init_command_line_parser.add_listener
def _(parser):
    parser.add_argument("--my-argument", type=str, env_var="LOCUST_MY_ARGUMENT", default="", help="它在工作")
    parser.add_argument("--env", choices=["dev", "staging", "prod"], default="dev", help="环境")
    parser.add_argument("--my-ui-invisible-argument", include_in_web_ui=False, default="我不可见")
    parser.add_argument("--my-ui-password-argument", is_secret=True, default="我是一个秘密")


    parser.add_argument("--my-ui-boolean-argument", default=True)
    parser.add_argument("--my-ui-required-argument", is_required=True, default="我是必需的")

@events.test_start.add_listener
def _(environment, **kw):
    print(f"自定义参数：{environment.parsed_options.my_argument}")

class WebsiteUser(HttpUser):
    @task
    def my_task(self):
        print(f"my_argument={self.environment.parsed_options.my_argument}")
        print(f"my_ui_invisible_argument={self.environment.parsed_options.my_ui_invisible_argument}")

在分布式运行 Locust 时，自定义参数会在测试开始时自动转发给工作节点（但不能在测试实际开始之前依赖转发的参数）。

测试数据管理

有多种方式可以将测试数据引入到你的测试中（毕竟，你的测试只是一个 Python 程序，它可以做任何 Python 能做的事情）。Locust 的事件让你可以精确控制何时获取/释放测试数据。你可以在这里找到详细的示例。