自定义负载形状

有时，可能需要一个完全自定义的负载测试形状，无法仅通过设置或更改用户数和启动速率来实现。例如，您可能希望在自定义时间生成负载峰值或进行上下波动。通过使用 LoadTestShape 类，您可以完全控制在任何时刻的用户数和启动速率。

定义一个继承自 LoadTestShape 类的类。如果在 locust 文件中找到了这种类型的类，Locust 将自动使用它。

在此类中，您定义一个 tick() 方法，该方法返回一个元组，包含所需的用户数和启动速率（或返回 None 以停止测试）。Locust 会大约每秒调用一次 tick() 方法。

您还可以在类中使用 get_run_time() 方法，来检查测试已经运行了多长时间。

示例

这个形状类会每 100 个用户块增加一次用户数，并且在 10 分钟后停止负载测试：

class MyCustomShape(LoadTestShape):
    time_limit = 600
    spawn_rate = 20

    def tick(self):
        run_time = self.get_run_time()

        if run_time < self.time_limit:
            # 用户数取最接近的百位数
            user_count = round(run_time, -2)
            return (user_count, self.spawn_rate)

        return None

这个功能可以在 GitHub 上的示例中进一步演示，包括：

  • 生成双波形（Double Wave Shape）

  • 基于时间的阶段，如 K6

  • 步骤负载模式，如 Visual Studio

还有一个可能对您的自定义负载形状有帮助的方法：get_current_user_count()，它返回当前活动用户的总数。这个方法可以用来防止在到达所需的用户数之前进入下一个步骤。如果每个用户的初始化过程较慢或不稳定，这个方法特别有用。如果这是您的使用案例，请参阅 GitHub 上的示例。

如果您想定义自己的自定义基础形状，则需要将抽象属性设置为 True，以避免它在导入时被当作形状使用：

class MyBaseShape(LoadTestShape):
    abstract = True

    def tick(self):
        # 一些可重用的功能，但需要继承
        return None

结合不同负载配置的用户

如果您使用 Web UI，您可以添加 --class-picker 参数来选择要使用的负载形状。但是，通常将用户定义放在一个文件中，将 LoadTestShape 放在另一个文件中会更加灵活。例如，如果您分别在 low_load.py 和 high_load.py 中定义了低负载和高负载形状：

locust -f locustfile.py,low_load.py

locust -f locustfile.py,high_load.py

限制每次 tick() 中生成的用户类型

通过向返回值中添加 user_classes 元素，您可以更精细地控制每次生成的用户类型：

class StagesShapeWithCustomUsers(LoadTestShape):

    stages = [
        {"duration": 10, "users": 10, "spawn_rate": 10, "user_classes": [UserA]},
        {"duration": 30, "users": 50, "spawn_rate": 10, "user_classes": [UserA, UserB]},
        {"duration": 60, "users": 100, "spawn_rate": 10, "user_classes": [UserB]},
        {"duration": 120, "users": 100, "spawn_rate": 10, "user_classes": [UserA, UserB]},
    ]

    def tick(self):
        run_time = self.get_run_time()

        for stage in self.stages:
            if run_time < stage["duration"]:
                try:
                    tick_data = (stage["users"], stage["spawn_rate"], stage["user_classes"])
                except:
                    tick_data = (stage["users"], stage["spawn_rate"])
                return tick_data

        return None

这个形状会在前 10 秒生成 10 个 UserA，接下来的 20 秒生成 40 个 UserA 和 UserB，直到阶段结束。

在自定义形状中重用通用选项

使用形状时，Web UI 中的用户数、启动速率和运行时间选项会被隐藏。如果您在命令行中指定它们，Locust 会记录警告。这是因为这些选项与形状不直接相关，指定它们可能会导致错误。

如果您确实希望将形状与这些选项结合使用，可以设置 use_common_options 属性，并从 self.runner.environment.parsed_options 中访问它们：

class MyCustomShape(LoadTestShape):
    use_common_options = True

    def tick(self):
        run_time = self.get_run_time()
        if run_time < self.runner.environment.parsed_options.run_time:
            # 用户数取最接近的百位数，就像在前面的示例中
            user_count = round(run_time, -2)
            return (user_count, self.runner.environment.parsed_options.spawn_rate)

        return None