4. 测试脚本编写：基于Python + Locust编写基础压测脚本

好，到了真正动手的环节了。前面我们聊了那么多理论、架构、指标，说白了都是为了这一刻——把脚本跑起来，看看系统到底扛不扛得住。

我个人习惯用Locust来做压测。为什么？因为它够轻量，写起来快，而且用Python写脚本，对我们做量化的人来说太友好了。你想想看，我们平时写策略、写回测框架，哪个不是Python？

4.1 为什么选Locust？

我记得第一次做交易系统压测时，团队里有人提议用JMeter。嗯，JMeter确实强大，但配置起来太繁琐了。尤其是模拟下单、撤单这种有状态的操作，你得搞一堆前置处理器、后置处理器，光看着就头大。

Locust就不一样了。它本质上就是一个Python库，你写一个类，继承HttpUser，然后在里面定义任务。每个任务就是一个普通的Python方法，想怎么控制逻辑就怎么控制。

核心优势：

  • 纯Python编写，零配置门槛
  • 支持协程，单机就能模拟大量用户
  • 内置Web UI，实时看结果
  • 任务之间可以共享状态，适合模拟交易流程

4.2 基础脚本结构

我们先搭一个最基础的架子。别急着写复杂逻辑，先把框架跑通再说。

from locust import HttpUser, task, between
import json

class TradingUser(HttpUser):
    # 每个用户执行任务之间的等待时间（模拟思考时间）
    wait_time = between(0.5, 2.0)
    
    def on_start(self):
        """用户启动时执行一次，用于登录"""
        # 我在项目中遇到过，如果不做登录态管理，
        # 压测出来的结果全是401，白忙活一场
        payload = {
            "username": "test_user",
            "password": "test_pass"
        }
        resp = self.client.post("/api/login", json=payload)
        if resp.status_code == 200:
            self.token = resp.json().get("token")
            self.client.headers.update({
                "Authorization": f"Bearer {self.token}"
            })
            print(f"用户登录成功，token: {self.token[:20]}...")
        else:
            print(f"登录失败: {resp.status_code}")
    
    @task(3)  # 权重为3，表示这个任务被选中的概率更高
    def query_balance(self):
        """查询账户余额"""
        resp = self.client.get("/api/account/balance")
        if resp.status_code != 200:
            print(f"查询余额失败: {resp.status_code}")
    
    @task(2)
    def query_positions(self):
        """查询持仓"""
        resp = self.client.get("/api/account/positions")
        # 这里可以加一些断言，但我建议压测时别加太多，
        # 否则会影响性能数据
        if resp.status_code == 200:
            positions = resp.json().get("positions", [])
            # 简单记录一下持仓数量
            if len(positions) > 0:
                pass  # 实际项目中可以记录到日志
    
    @task(1)
    def place_order(self):
        """模拟下单"""
        order = {
            "symbol": "BTCUSDT",
            "side": "buy",
            "type": "limit",
            "price": 50000.0,
            "quantity": 0.01
        }
        resp = self.client.post("/api/order", json=order)
        if resp.status_code == 200:
            order_id = resp.json().get("order_id")
            # 把订单ID存起来，后面撤单要用
            if not hasattr(self, "order_ids"):
                self.order_ids = []
            self.order_ids.append(order_id)
            # 只保留最近10个订单，防止内存泄漏
            if len(self.order_ids) > 10:
                self.order_ids.pop(0)
    
    @task(1)
    def cancel_order(self):
        """模拟撤单"""
        if not hasattr(self, "order_ids") or len(self.order_ids) == 0:
            # 没有可撤的订单，就跳过
            return
        order_id = self.order_ids.pop(0)
        resp = self.client.delete(f"/api/order/{order_id}")
        if resp.status_code != 200:
            print(f"撤单失败: {order_id}, 状态码: {resp.status_code}")

小技巧：我在做压测时，习惯把每个任务的权重调成跟真实用户行为一致。比如真实用户可能80%的时间在查询，10%在下单，10%在撤单。那权重就设成8:1:1。这样压出来的数据才有参考价值。

4.3 模拟真实交易流程

上面那个脚本虽然能跑，但太简单了。真实交易场景下，用户的行为是有顺序的：先查行情，再下单，然后查订单状态，最后可能撤单或者等成交。

我建议用@task装饰器配合SequentialTaskSet来实现流程控制。

from locust import HttpUser, task, between, SequentialTaskSet

class TradingFlow(SequentialTaskSet):
    """模拟一个完整的交易流程"""
    
    def on_start(self):
        self.order_id = None
        self.symbol = "BTCUSDT"
    
    @task
    def step1_query_market(self):
        """第一步：查询行情"""
        resp = self.client.get(f"/api/market/ticker?symbol={self.symbol}")
        if resp.status_code == 200:
            ticker = resp.json()
            self.current_price = float(ticker.get("lastPrice", 50000))
            print(f"当前价格: {self.current_price}")
    
    @task
    def step2_place_order(self):
        """第二步：下单"""
        order = {
            "symbol": self.symbol,
            "side": "buy",
            "type": "limit",
            "price": self.current_price * 0.99,  # 比当前价低1%
            "quantity": 0.01
        }
        resp = self.client.post("/api/order", json=order)
        if resp.status_code == 200:
            self.order_id = resp.json().get("order_id")
            print(f"下单成功，订单ID: {self.order_id}")
    
    @task
    def step3_query_order(self):
        """第三步：查询订单状态"""
        if self.order_id is None:
            return
        resp = self.client.get(f"/api/order/{self.order_id}")
        if resp.status_code == 200:
            status = resp.json().get("status")
            print(f"订单状态: {status}")
            # 如果已经成交，就不撤单了
            if status == "filled":
                self.order_id = None
    
    @task
    def step4_cancel_order(self):
        """第四步：撤单（如果还没成交）"""
        if self.order_id is None:
            return
        resp = self.client.delete(f"/api/order/{self.order_id}")
        if resp.status_code == 200:
            print(f"撤单成功: {self.order_id}")
            self.order_id = None

class TradingUser(HttpUser):
    wait_time = between(1, 3)
    tasks = [TradingFlow]

注意：我曾经犯过一个错误——在压测脚本里加了大量的print语句。结果压测还没跑起来，控制台先被日志淹没了，CPU直接飙到100%。压测脚本本身也会消耗资源，记得把日志级别调高，或者只在调试时打开。

4.4 参数化与数据隔离

每个虚拟用户应该用不同的数据，不然就变成所有人在抢同一个订单，压出来的结果会失真。

我一般会在on_start里给每个用户分配一个独立的账户。

import random
import string

class TradingUser(HttpUser):
    wait_time = between(0.5, 2.0)
    
    def on_start(self):
        # 生成随机用户名，确保每个用户独立
        suffix = ''.join(random.choices(string.ascii_lowercase, k=8))
        self.username = f"load_test_{suffix}"
        
        # 登录
        payload = {
            "username": self.username,
            "password": "test123"
        }
        resp = self.client.post("/api/login", json=payload)
        if resp.status_code == 200:
            self.token = resp.json().get("token")
            self.client.headers.update({
                "Authorization": f"Bearer {self.token}"
            })
        
        # 初始化账户资金（如果有初始化接口的话）
        self.client.post("/api/account/init", json={
            "balance": 100000.0  # 每个用户10万初始资金
        })
    
    @task
    def random_order(self):
        """随机下单，模拟不同用户的交易行为"""
        symbols = ["BTCUSDT", "ETHUSDT", "BNBUSDT"]
        sides = ["buy", "sell"]
        
        order = {
            "symbol": random.choice(symbols),
            "side": random.choice(sides),
            "type": "limit",
            "price": round(random.uniform(40000, 60000), 2),
            "quantity": round(random.uniform(0.001, 0.1), 4)
        }
        self.client.post("/api/order", json=order)

4.5 运行与监控

脚本写好了，怎么跑？我推荐两种方式：

方式 命令 适用场景
Web UI模式 locust -f trading_test.py --web-host 0.0.0.0 调试脚本、小规模压测
无头模式 locust -f trading_test.py --headless -u 100 -r 10 --run-time 10m 自动化压测、CI/CD集成

无头模式参数说明：

  • -u 100：模拟100个并发用户
  • -r 10：每秒启动10个用户（爬坡速率）
  • --run-time 10m：运行10分钟

我的经验：刚开始压测时，别一上来就上1000个用户。先跑10个，看看系统响应时间。没问题了再慢慢往上加。我见过太多人一上来就把系统压崩了，结果连问题出在哪都不知道。

4.6 避坑指南

嗯，最后分享几个我踩过的坑：

  • 连接池耗尽：Locust默认每个用户一个连接，但如果你在任务里频繁创建新的HTTP会话，连接池会很快耗尽。解决办法是用self.client，它内部已经做了连接复用。
  • 状态污染：不同用户之间不要共享全局变量。我见过有人用全局列表存订单ID，结果所有用户都在操作同一个列表，数据全乱了。
  • 超时设置：交易系统一般对延迟要求很高。建议在self.client里设置超时时间，比如self.client.timeout = 5。如果5秒没响应，直接算失败，别傻等。
  • 数据清理：压测结束后，记得清理测试数据。我曾经忘了清理，结果生产环境里多了几万条测试订单，被运维骂了一顿。

脚本写好了，下一步就是跑起来看结果了。不过那是下一节的内容，我们先把脚本调通再说。

专注资料整理

← 上一章 📖 返回目录 下一章 →