【性能测试利器】使用Locust：一个面向开发者的分布式负载测试框架

【性能测试利器】使用Locust：一个面向开发者的分布式负载测试框架

1. 引言：为什么需要Locust？

在AI驱动的现代应用架构中，后端服务（如模型推理API、特征提取服务、数据查询接口）的性能和稳定性直接决定了用户体验和系统可靠性。传统的压力测试工具如JMeter虽然功能强大，但其配置繁琐、基于UI的模式难以融入CI/CD pipeline，并且对于复杂的、动态变化的测试场景支持不足。

Locust应运而生。它是一个用Python编写的开源负载测试工具，其核心理念是**“用代码定义你的测试”**。这带来了无与伦比的灵活性和可扩展性：

• 开发者友好： 你可以用纯Python代码模拟复杂的用户流程、动态参数化和自定义校验逻辑，无缝集成到你的开发环境中。
• 分布式与可扩展： 单机可以模拟数千用户，通过主从架构轻松实现横向扩展，以满足百万级并发的要求。
• 实时Web UI： 无需复杂的配置，启动后即可通过浏览器实时查看测试指标（RPS、响应时间、失败率等），并动态调整负载。
• AI场景契合： 非常适合测试需要传递复杂JSON payload（如模型输入）、处理长连接（如WebSocket）或具有特定认证机制（如JWT）的AI服务。

本文将带你深入探索Locust的核心功能，并通过具体的代码示例展示其强大之处。

---

2. Locust核心功能与使用案例

2.1. 核心概念：User, Task, Environment

在编写Locust脚本时，你需要理解三个核心类：

• HttpUser： 代表一类用户（或一个用户组）。你可以在其中定义这类用户的等待时间、权重以及他们将要执行的任务集合。
• TaskSet： 定义一系列任务的集合，可以嵌套，用于模拟用户复杂的操作序列。
• task： 一个装饰器，用于将方法声明为用户的一个可执行任务。你可以为任务设置权重，权重越高，被执行的频率越高。

2.2. 功能一：基础HTTP请求测试

这是Locust最基础的功能，用于测试简单的RESTful API。

示例：测试一个API网关后的文本分类模型服务

from locust import HttpUser, task, between

class AITestUser(HttpUser):
    # 模拟用户在每个任务执行后等待1-3秒
    wait_time = between(1, 3)

    @task
    def classify_text(self):
        # 定义请求头，表明传递的是JSON数据
        headers = {"Content-Type": "application/json"}
        # 定义请求体，模拟模型输入
        payload = {
            "text": "This is a fantastic movie! Great acting and plot."
        }
        
        # 使用client发起POST请求，client是HttpUser自带的属性，基于requests库
        with self.client.post(
            "/v1/models/classify:predict",
            json=payload,
            headers=headers,
            catch_response=True # 允许手动控制成功/失败判断
        ) as response:
            # 手动校验响应：状态码为200且结果中包含"positive"才算成功
            if response.status_code == 200:
                if "positive" in response.text:
                    response.success()
                else:
                    response.failure(f"Unexpected result: {response.text}")
            else:
                response.failure(f"Status code was {response.status_code}")

if __name__ == "__main__":
    # 通常通过命令行运行，但也可以直接运行脚本
    import os
    os.system("locust -f this_script.py")

代码解释：

1. 我们创建了一个 AITestUser 类，它继承自 HttpUser。
2. wait_time = between(1, 3) 定义了用户行为间隔，使负载更接近真实用户。
3. @task 装饰器将 classify_text 方法标记为一个任务。
4. 使用 self.client 发送HTTP POST请求到模型的预测端点。
5. catch_response=True 和 with 语句块允许我们根据响应内容（而不仅仅是HTTP状态码）来手动标记请求的成功或失败。这对于测试AI API至关重要，因为返回200但结果错误依然是失败。

2.3. 功能二：任务权重与嵌套

你可以控制不同任务的执行比例，并使用嵌套的TaskSet来组织复杂的用户行为流。

示例：模拟一个混合了“浏览”和“购买”行为的电商用户

from locust import HttpUser, task, TaskSet, between

class BrowseProducts(TaskSet):
    # 这个任务的权重是3，意味着执行频率是`view_cart`的3倍
    @task(3)
    def view_products(self):
        self.client.get("/products")
        
    @task(1)
    def view_cart(self):
        self.client.get("/cart")

    # 用户可以中断当前TaskSet，跳回到父User中定义的任务
    @task(1)
    def stop(self):
        self.interrupt()

class PurchaseFlow(TaskSet):
    # 在TaskSet的on_start方法中定义的任务，每个用户实例只会执行一次
    def on_start(self):
        self.client.post("/login", json={"username": "test_user", "password": "secret"})

    @task
    def checkout(self):
        self.client.post("/checkout", json={"item_id": 123})

    @task(1)
    def stop(self):
        self.interrupt()

class WebsiteUser(HttpUser):
    wait_time = between(2, 5)
    # tasks是一个列表，可以接收TaskSet类或(权重, TaskSet)元组
    tasks = {
        BrowseProducts: 4,  # 权重4，用户有更高概率执行浏览任务
        PurchaseFlow: 1     # 权重1，用户有较低概率执行购买流程
    }

    # 在HttpUser层面也可以定义任务
    @task
    def index(self):
        self.client.get("/")

代码解释：

1. 我们定义了两个 TaskSet：BrowseProducts（浏览行为）和 PurchaseFlow（购买行为）。
2. 在 WebsiteUser 中，使用 tasks 列表以权重的形式分配这两个TaskSet，4:1的比例表示用户大部分时间在浏览，偶尔会购买。
3. PurchaseFlow 中的 on_start 方法是一个特殊方法，每个模拟用户在执行这个TaskSet前都会先执行一次，常用于登录等初始化操作。
4. self.interrupt() 是退出当前嵌套TaskSet的关键，允许控制流返回到父级。

2.4. 功能三：参数化与数据驱动测试

真实的测试场景需要不同的用户使用不同的数据。Locust可以轻松地从外部文件（如CSV、JSON）中读取数据以实现参数化。

示例：模拟多个用户使用不同的用户名和测试文本进行请求

首先，创建一个 users.csv 文件：

username,test_text
user1,This is amazing.
user2,I'm not sure about this.
user3,Total waste of time.

Locust脚本：

from locust import HttpUser, task, between
import csv
import random

class ParameterizedUser(HttpUser):
    wait_time = between(1, 2.5)

    def on_start(self):
        # 在用户启动时，从CSV文件中读取所有数据
        self.data = []
        with open('users.csv', 'r') as f:
            reader = csv.DictReader(f)
            for row in reader:
                self.data.append(row)
        
    @task
    def sentimen_analysis(self):
        # 随机从数据列表中选取一条记录
        if self.data:
            user_data = random.choice(self.data)
            payload = {
                "username": user_data['username'],
                "text": user_data['test_text']
            }
            with self.client.post("/analyze", json=payload, catch_response=True) as resp:
                # ... 处理响应 ...
                pass

代码解释：

1. on_start 在每个模拟用户开始运行时被调用，在这里我们读取CSV文件并将数据存储在实例变量 self.data 中。
2. 在执行 sentimen_analysis 任务时，使用 random.choice 从数据列表中随机选取一条记录，用于构造请求的payload。
3. 这种方法确保了并发用户会使用不同的测试数据，避免了所有请求都完全一样的“脱敏”测试，使得测试结果更具代表性。

2.5. 功能四：自定义客户端（测试非HTTP协议）

Locust的强大之处在于其可扩展性。通过继承 User 类而非 HttpUser，你可以使用任何Python库来测试任何协议，如WebSocket、MQTT、gRPC等。

示例：测试一个WebSocket服务（例如实时AI通知服务）

你需要先安装 websocket-client 库。

pip install websocket-client

Locust脚本：

from locust import User, task, between
import websocket
import json
import time

class WebSocketClient:
    def __init__(self, host):
        # 创建WebSocket连接
        self.ws = websocket.create_connection(f"ws://{host}/ws/notifications/")
        
    def send(self, payload):
        # 发送消息
        start_time = time.time()
        self.ws.send(json.dumps(payload))
        # 接收回声或响应（简单示例）
        result = self.ws.recv()
        total_time = int((time.time() - start_time) * 1000) # 计算耗时(ms)
        # 这里可以记录成功/失败，为了示例简单，我们总是记录成功
        return total_time

    def on_close(self):
        self.ws.close()

class WebSocketUser(User):
    wait_time = between(1, 3)
    
    def on_start(self):
        # 在用户启动时初始化自定义客户端
        self.client = WebSocketClient(self.host)
        
    @task
    def send_notification(self):
        payload = {"user_id": 123, "message": "AI task completed"}
        response_time = self.client.send(payload)
        # 手动触发Locust的事件记录
        self.environment.events.request.fire(
            request_type="WS",
            name="Send Notification",
            response_time=response_time,
            response_length=len(str(payload)),
            exception=None,
        )
        
    def on_stop(self):
        # 在用户停止时关闭连接
        self.client.on_close()

代码解释：

1. 我们创建了一个自定义的 WebSocketClient 类来处理WebSocket连接和消息收发。
2. WebSocketUser 继承自基础的 User 类。在 on_start 中初始化我们的客户端。
3. 在任务中，我们调用 self.client.send() 并手动使用 self.environment.events.request.fire() 来向Locust报告请求的耗时和状态。这是将自定义协议测试集成到Locust统计中的关键。
4. 这种方式使得Locust几乎可以测试任何你能用Python编写的协议。

---

3. 总结

Locust以其“代码即一切”的理念，为开发者提供了一个极其强大且灵活的性能测试工具。通过本文的探讨，我们可以总结出它的几个关键优势：

1. 灵活性与可读性： 用Python代码定义测试，逻辑清晰，易于维护和版本控制，非常适合融入CI/CD流程。
2. 强大的场景模拟能力： 通过任务权重、嵌套TaskSet和参数化，能够构建出高度模拟真实世界用户行为的复杂测试场景。
3. 卓越的可扩展性： 不局限于HTTP，可以轻松扩展至WebSocket、gRPC、MQTT等任何协议，满足AI和IoT等前沿领域的测试需求。
4. 直观的实时监控： 内置的Web UI提供了丰富的实时指标，帮助你快速识别性能瓶颈。

无论你是在测试一个传统的Web应用，还是一个复杂的、由AI驱动的微服务生态系统，Locust都提供了一个现代化、可扩展且高效的解决方案。告别笨重的UI配置，拥抱代码，让你的性能测试变得真正强大而优雅。