如何在Locust的LoadTestShape类__init__方法中获取CLI参数

根因说明 #

Locust自动实例化LoadTestShape子类的过程中，会先执行子类的__init__方法，再注入environment、runner属性，所以原生初始化逻辑里拿不到CLI参数。

最优方案：手动实例化LoadShape（符合OOP设计） #

步骤1：注册自定义CLI参数 #

首先在locustfile中注册你需要的命令行参数：

from locust import LoadTestShape, events
from locust.argument_parser import LocustArgumentParser

# 注册自定义CLI参数
@events.init_command_line_parser.add_listener
def add_custom_args(parser: LocustArgumentParser):
    parser.add_argument("--max-rps", type=int, default=500, help="压测最大RPS")
    parser.add_argument("--step-interval", type=int, default=5, help="阶梯升压间隔时长（秒）")

步骤2：改造自定义LoadShape的构造函数 #

调整MyCustomLoadShape的__init__方法，支持直接传入配置参数：

class CustomLoadShape:
    # 工厂类保留原有逻辑
    def __init__(self, step_interval=5, max_rps=500):
        self.step_interval = step_interval
        self.max_rps = max_rps

    def _tick_step_increase(self):
        pass

    def _tick_square_wave(self):
        pass

    def _tick_rps_based(self):
        # 你的原有tick逻辑
        pass

class MyCustomLoadShape(LoadTestShape, CustomLoadShape):
    # 改造构造函数支持参数传入
    def __init__(self, step_interval=5, max_rps=500, *args, **kwargs):
        super().__init__(*args, **kwargs)
        CustomLoadShape.__init__(self, step_interval=step_interval, max_rps=max_rps)
        self.tick = self._tick_rps_based

步骤3：在初始化事件中手动实例化LoadShape #

通过Locust的init事件，在环境初始化完成后，手动创建LoadShape实例并传入CLI参数，替换默认的自动实例化对象：

@events.init.add_listener
def on_locust_init(environment, **kwargs):
    # 从已解析的CLI参数中取值
    max_rps = environment.parsed_options.max_rps
    step_interval = environment.parsed_options.step_interval
    # 手动实例化自定义LoadShape，赋值给environment替换默认对象
    environment.load_shape = MyCustomLoadShape(
        step_interval=step_interval,
        max_rps=max_rps
    )

轻量替代方案：延迟初始化 #

如果不想改动事件逻辑，可以在LoadShape中加初始化标记，第一次执行tick时完成属性配置，仅执行一次：

class MyCustomLoadShape(LoadTestShape, CustomLoadShape):
    def __init__(self):
        self.initialized = False
        self.tick = self._tick_rps_based

    def _lazy_init(self):
        if self.initialized:
            return
        # 第一次执行tick时runner已注入，可正常获取参数
        max_rps = self.runner.environment.parsed_options.max_rps
        step_interval = self.runner.environment.parsed_options.step_interval
        CustomLoadShape.__init__(self, step_interval=step_interval, max_rps=max_rps)
        self.initialized = True

    def _tick_rps_based(self):
        self._lazy_init()
        # 原有tick逻辑

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