如何将Locust的StagesShape动态参数传递至UserInstance类

方法1：使用全局变量共享参数 #

直接通过模块级全局变量存储当前阶段的param1、param2，User实例发起请求时读取该变量即可。

# 全局变量，存储当前阶段的请求头参数
current_stage_headers = {"param1": "", "param2": ""}

class CustomStagesShape(StagesShape):
    def __init__(self):
        # 读取CSV并解析测试阶段
        self.stages = self._load_stages_from_csv("test_config.csv")
        self.current_stage_index = 0

    def _load_stages_from_csv(self, csv_path):
        stages = []
        with open(csv_path, "r", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            for row in reader:
                # 转换datetime为时间戳，方便阶段判断
                end_timestamp = datetime.strptime(row["datetime"], "%Y-%m-%d %H:%M:%S").timestamp()
                stages.append({
                    "end_time": end_timestamp,
                    "rps": int(row["rps"]),
                    "param1": row["param1"],
                    "param2": row["param2"]
                })
        return stages

    def tick(self):
        current_time = time.time()
        if self.current_stage_index < len(self.stages):
            current_stage = self.stages[self.current_stage_index]
            if current_time < current_stage["end_time"]:
                # 更新全局参数
                global current_stage_headers
                current_stage_headers["param1"] = current_stage["param1"]
                current_stage_headers["param2"] = current_stage["param2"]
                # 返回当前用户数和RPS（可根据需求调整用户数计算逻辑）
                return (int(current_stage["rps"] * 1.5), current_stage["rps"])
            else:
                self.current_stage_index += 1
        # 所有阶段执行完毕，返回None停止测试
        return None

class UserInstance(HttpUser):
    @task
    def execute_post_request(self):
        # 从全局变量获取当前阶段的请求头参数
        request_headers = {
            "param1": current_stage_headers["param1"],
            "param2": current_stage_headers["param2"]
        }
        # 发起POST请求
        self.client.post("/target-api", headers=request_headers)

方法2：通过Environment实例共享参数 #

利用Locust内置的Environment实例共享数据，该方式支持分布式测试场景下的参数自动同步。

class CustomStagesShape(StagesShape):
    def __init__(self, environment):
        super().__init__(environment)
        self.stages = self._load_stages_from_csv("test_config.csv")
        self.current_stage_index = 0
        # 初始化Environment的user_data存储参数
        self.environment.user_data["param1"] = ""
        self.environment.user_data["param2"] = ""

    def _load_stages_from_csv(self, csv_path):
        stages = []
        with open(csv_path, "r", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            for row in reader:
                end_timestamp = datetime.strptime(row["datetime"], "%Y-%m-%d %H:%M:%S").timestamp()
                stages.append({
                    "end_time": end_timestamp,
                    "rps": int(row["rps"]),
                    "param1": row["param1"],
                    "param2": row["param2"]
                })
        return stages

    def tick(self):
        current_time = time.time()
        if self.current_stage_index < len(self.stages):
            current_stage = self.stages[self.current_stage_index]
            if current_time < current_stage["end_time"]:
                # 更新Environment中的参数
                self.environment.user_data["param1"] = current_stage["param1"]
                self.environment.user_data["param2"] = current_stage["param2"]
                return (int(current_stage["rps"] * 1.5), current_stage["rps"])
            else:
                self.current_stage_index += 1
        return None

class UserInstance(HttpUser):
    @task
    def execute_post_request(self):
        # 从Environment获取当前阶段的参数
        request_headers = {
            "param1": self.environment.user_data["param1"],
            "param2": self.environment.user_data["param2"]
        }
        self.client.post("/target-api", headers=request_headers)

关键注意事项 #

• 全局变量方法适合单机测试，分布式场景下主节点的全局变量无法自动同步到工作节点，此时优先使用Environment方案
• Environment.user_data会自动在分布式模式下同步到所有工作节点，无需额外处理
• 确保CSV中的datetime格式与代码中的解析格式完全匹配，否则会导致阶段切换逻辑失效

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