如何从CSV文件读取参数替代Locust脚本中硬编码的TSIDs？

实现步骤 #

1. 准备CSV文件
   在脚本同目录下创建tsids.csv，格式如下（每行一个TSID，第一行可加表头）：

tsid
128132
85624
123456
789012

2. 修改Locust脚本
   用Python内置的csv模块读取CSV文件，全局加载TSIDs数据（避免每个用户实例重复读取文件），同时修正原脚本中的语法错误。

修改后的完整脚本：

from locust import HttpUser, task
import random
import csv
import os
from datetime import datetime, timedelta

server_name = "https://xxx"

# 从同目录CSV读取TSIDs
TSIDs = set()
# 获取脚本所在目录，确保路径可靠
csv_path = os.path.join(os.path.dirname(__file__), "tsids.csv")
with open(csv_path, mode='r', newline='', encoding='utf-8') as f:
    reader = csv.DictReader(f)
    for row in reader:
        TSIDs.add(row['tsid'].strip())

def set_from_date():
    # 示例：设置为当前时间前7天的UTC时间
    from_date = (datetime.utcnow() - timedelta(days=7)).isoformat() + "Z"
    return from_date

def set_to_date():
    # 示例：设置为当前UTC时间
    to_date = datetime.utcnow().isoformat() + "Z"
    return to_date

class LoadValues(HttpUser):
    host = server_name

    def _run_read_ts(self, series_list, resolution):
        ts_ids_str = ",".join(series_list)
        # 修正原脚本endUtc参数缺失=的问题
        resp = self.client.get(
            f'/api/loadValues?tsIds={ts_ids_str}&resolution={resolution}&startUtc={set_from_date()}&endUtc={set_to_date()}',
            headers={'X-API-KEY': 'AKjCg9hTcYQ='}
        )
        print("Response status code:", resp.status_code)

    @task(1)
    def test_get_ts_1(self):
        # 从读取的TSIDs中随机选1个
        self._run_read_ts(random.sample(TSIDs, 1), 'PT15M')

关键说明 #

• 全局加载数据：在脚本开头读取CSV并加载到TSIDs集合，所有用户实例共享这份数据，避免频繁IO操作影响测试性能。
• 路径可靠性：用os.path.dirname(__file__)获取脚本所在目录，再拼接CSV文件名，不管脚本从哪个目录执行都能正确找到文件。
• 数据格式兼容：保持TSIDs为集合类型，和原脚本的random.sample用法兼容，确保随机选值逻辑不变。
• 语法修正：原脚本中@task缩进错误、endUtc参数后缺少=符号，已在修改后的代码中修复，同时简化了_run_read_ts的参数（无需重复传入start/end，直接调用函数获取）。

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