如何通过命令行为Locust指定自定义负载形状类

默认Locust要求把自定义负载形状类直接写在存放用户类的locust文件里，切换负载逻辑必须改代码，灵活度很差，要实现命令行动态指定，按下面的方法操作即可。

首先明确Locust自定义负载的核心规则：

当需要实现突增突降、分阶段升降压这类没法靠固定用户数、固定孵化速率实现的负载场景时，可以继承LoadTestShape实现完全自定义的负载逻辑：

• 类里必须实现tick()方法，Locust每秒会调用一次这个方法，返回(目标用户数, 每秒孵化速率)格式的元组，返回None就会自动终止测试
• 类内部可以调用get_run_time()方法，获取测试已经运行的时长（单位：秒），用来做时间维度的负载判断
• 只要Locust运行环境里绑定了合法的LoadTestShape子类实例，就会自动启用自定义负载模式，忽略命令行传入的固定用户数、孵化速率参数

具体操作步骤：

• 把所有自定义负载形状抽到独立Python文件里维护，不要硬写在主locustfile中。比如新建shapes.py，在里面写所有负载形状的实现类，示例：

from locust import LoadTestShape

# 突增负载示例
class SpikeShape(LoadTestShape):
    def tick(self):
        run_time = self.get_run_time()
        if run_time < 30:
            return (10, 2)
        elif 30 <= run_time < 60:
            return (1000, 100)
        elif 60 <= run_time < 90:
            return (10, 2)
        else:
            return None

# 阶梯加压示例
class StepLoadShape(LoadTestShape):
    def tick(self):
        run_time = self.get_run_time()
        current_step = run_time // 60
        target_user = min(current_step * 50, 300)
        if target_user >= 300 and run_time > 360:
            return None
        return (target_user, 10)

• 在主locustfile.py里添加自定义命令行参数，测试启动时根据传入的参数动态加载对应的负载形状类，绑定到Locust运行环境即可，示例代码：

from locust import HttpUser, task, events
import importlib

# 这里写你自己的压测用户逻辑
class DemoUser(HttpUser):
    @task
    def visit_home(self):
        self.client.get("/")

# 注册自定义命令行参数
@events.init_command_line_parser.add_listener
def _(parser):
    parser.add_argument(
        "--load-shape",
        type=str,
        default="",
        help="指定自定义负载形状，格式为 模块名:类名，比如 shapes:SpikeShape"
    )

# 测试启动时动态加载指定的负载形状
@events.test_start.add_listener
def _(environment, **kwargs):
    shape_arg = environment.parsed_options.load_shape
    if not shape_arg:
        return
    module_name, cls_name = shape_arg.split(":")
    mod = importlib.import_module(module_name)
    shape_cls = getattr(mod, cls_name)
    environment.shape_class = shape_cls()

• 启动测试时直接通过--load-shape参数指定要用的负载形状即可，不需要修改locustfile代码：
  • 启动突增负载测试：locust -f locustfile.py --load-shape shapes:SpikeShape
  • 启动阶梯加压测试：locust -f locustfile.py --load-shape shapes:StepLoadShape

内容的提问来源于stack exchange，提问作者Alexandre V.