使用--class-picker时Locust UI仅加载单个LoadTestShape类的问题

问题原因 #

Locust的--class-picker参数仅用于选择要运行的User类，对LoadTestShape类不生效。根据Locust默认逻辑，当文件中定义多个LoadTestShape子类时，它只会自动加载并使用第一个被定义的类，不会在Web UI中列出所有Shape类供选择。

解决方案 #

方案1：通过命令行指定要使用的Shape类

使用--shape-class参数直接指定要加载的LoadTestShape子类名称，启动命令示例：

• 使用第一个Shape类：

locust -f locustfile.py --shape-class ramp_up_1min_40users_hold_20min

• 使用第二个Shape类：

locust -f locustfile.py --shape-class ramp_up_20min_40users_hold_5min_ramp_down_5min_1user

方案2：实现动态可配置的Shape类

如果需要更灵活的切换方式，可以将所有负载阶段配置整合到一个Shape类中，通过环境变量指定要使用的配置：

修改后的代码示例：

import os
from locust import HttpUser, TaskSet, task, LoadTestShape

class Tasks1(TaskSet):
    @task
    def service(self):
        print("Launching URL")

    @task
    def stop(self):
        self.interrupt()

class FirstUser(HttpUser):
    tasks = [Tasks1]

# 统一管理所有负载阶段配置
SHAPE_CONFIGS = {
    "short_ramp_hold": [
        {"duration": 60, "users": 40, "spawn_rate": 0.67},
        {"duration": 1200, "users": 40, "spawn_rate": 1},
    ],
    "long_ramp_rampdown": [
        {"duration": 1200, "users": 40, "spawn_rate": 0.034},
        {"duration": 300, "users": 40, "spawn_rate": 1},
        {"duration": 300, "users": 1, "spawn_rate": 0.0325}
    ]
}

class DynamicLoadTestShape(LoadTestShape):
    def __init__(self):
        super().__init__()
        # 从环境变量读取配置名，默认使用第一个配置
        config_name = os.getenv("LOCUST_SHAPE", list(SHAPE_CONFIGS.keys())[0])
        self.stages = SHAPE_CONFIGS[config_name]

    def tick(self):
        run_time = self.get_run_time()
        for stage in self.stages:
            if run_time < stage["duration"]:
                return (stage["users"], stage["spawn_rate"])
        return None

启动时通过环境变量指定要使用的配置：

LOCUST_SHAPE=long_ramp_rampdown locust -f locustfile.py

内容的提问来源于stack exchange，提问作者sheragon