Locust 负载测试框架详解

一、什么是 Locust？

Locust 是一个开源的负载测试工具，用于测试网站、API 或其他系统的性能和可扩展性。它使用 Python 编写，通过编写简单的 Python 代码来定义用户行为。

主要特点：

• 完全基于代码：测试场景用 Python 编写，灵活性高
• 分布式支持：可轻松扩展到多台机器
• 实时 Web UI：监控测试过程
• 可扩展：支持自定义客户端和协议
• 统计信息丰富：提供详细的性能指标

二、安装 Locust

# 基本安装
pip install locust

# 安装更多功能
pip install locust[extra]

三、基本架构

1. 核心组件

• User 类：定义虚拟用户行为
• TaskSet 类：定义任务集合
• HttpUser：内置的 HTTP 用户类
• Web UI：监控界面
• Master/Worker 模式：分布式测试架构

四、编写第一个 Locust 测试脚本

基本示例

from locust import HttpUser, TaskSet, task, between

# 定义任务集
class UserBehavior(TaskSet):
    @task(1)  # 任务权重，数值越大执行频率越高
    def view_homepage(self):
        self.client.get("/")
    
    @task(2)
    def view_products(self):
        self.client.get("/products")
    
    @task(1)
    def login(self):
        self.client.post("/login", {
            "username": "test",
            "password": "test123"
        })

# 定义用户类
class WebsiteUser(HttpUser):
    tasks = [UserBehavior]  # 关联任务集
    wait_time = between(1, 5)  # 等待时间范围（秒）
    host = "http://example.com"  # 目标主机

五、详细的测试脚本编写

1. 用户类和等待时间

from locust import HttpUser, task, between, constant, constant_pacing

class TestUser(HttpUser):
    # 等待时间策略：
    # between(min, max) - 随机等待
    wait_time = between(1, 3)
    
    # constant(seconds) - 固定间隔
    # wait_time = constant(2)
    
    # constant_pacing(seconds) - 保证任务执行间隔
    # wait_time = constant_pacing(2)
    
    host = "http://api.example.com"

2. 任务定义方式

from locust import HttpUser, task

# 方式1：使用 @task 装饰器
class ApiUser(HttpUser):
    @task(3)  # 权重为3
    def get_users(self):
        self.client.get("/api/users")
    
    @task(1)  # 权重为1
    def create_user(self):
        self.client.post("/api/users", json={
            "name": "John",
            "email": "john@example.com"
        })

# 方式2：使用 tasks 属性
class MixedTasksUser(HttpUser):
    def low_priority_task(self):
        self.client.get("/api/status")
    
    def high_priority_task(self):
        self.client.post("/api/data", json={"data": "test"})
    
    # 直接指定任务列表
    tasks = {
        low_priority_task: 1,
        high_priority_task: 3
    }

# 方式3：任务序列
class SequentialUser(HttpUser):
    @task
    def sequential_tasks(self):
        # 顺序执行一系列请求
        self.client.get("/step1")
        self.client.get("/step2")
        self.client.post("/step3", json={"action": "complete"})

3. 使用 TaskSet 组织复杂任务

from locust import HttpUser, TaskSet, task

class BrowseProducts(TaskSet):
    @task(10)
    def view_products(self):
        self.client.get("/products")
    
    @task(2)
    def view_product_detail(self):
        product_id = self.user.product_ids[0]  # 访问父用户属性
        self.client.get(f"/products/{product_id}")
    
    @task(1)
    def stop(self):
        self.interrupt()  # 退出当前 TaskSet

class AddToCart(TaskSet):
    @task
    def add_item(self):
        self.client.post("/cart/add", json={
            "product_id": "123",
            "quantity": 1
        })
    
    @task
    def view_cart(self):
        self.client.get("/cart")
    
    @task
    def checkout(self):
        self.client.post("/checkout")

class ECommerceUser(HttpUser):
    tasks = [BrowseProducts, AddToCart]
    product_ids = ["1001", "1002", "1003"]
    
    def on_start(self):
        """用户启动时执行"""
        self.login()
    
    def on_stop(self):
        """用户停止时执行"""
        self.logout()
    
    def login(self):
        response = self.client.post("/login", json={
            "username": "test",
            "password": "test123"
        })
        self.token = response.json().get("token")
    
    def logout(self):
        self.client.post("/logout", headers={
            "Authorization": f"Bearer {self.token}"
        })

六、高级功能

1. 自定义客户端

from locust import User, task, events
import websocket
import json

class WebSocketClient:
    def __init__(self, host):
        self.host = host
        self.ws = None
    
    def connect(self):
        self.ws = websocket.create_connection(self.host)
    
    def send(self, message):
        self.ws.send(json.dumps(message))
    
    def receive(self):
        return json.loads(self.ws.recv())
    
    def close(self):
        if self.ws:
            self.ws.close()

class WebSocketUser(User):
    abstract = True
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.client = WebSocketClient(self.host)
    
    def on_start(self):
        self.client.connect()
    
    def on_stop(self):
        self.client.close()

class MyWebSocketUser(WebSocketUser):
    @task
    def send_message(self):
        self.client.send({"type": "ping"})
        response = self.client.receive()
        print(f"Received: {response}")
    
    wait_time = between(1, 3)
    host = "ws://echo.websocket.org"

2. 事件钩子

from locust import events
from locust.runners import MasterRunner

# 测试开始事件
@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    print("测试开始")
    if isinstance(environment.runner, MasterRunner):
        print("这是 Master 节点")

# 测试停止事件
@events.test_stop.add_listener
def on_test_stop(environment, **kwargs):
    print("测试结束")

# 请求事件
@events.request.add_listener
def on_request(request_type, name, response_time, response_length, exception, context, **kwargs):
    if exception:
        print(f"请求失败: {name}, 异常: {exception}")
    else:
        print(f"请求成功: {name}, 响应时间: {response_time}ms")

# 自定义统计
from locust import stats
stats.PERCENTILES_TO_REPORT = [0.50, 0.95, 0.99]

3. 自定义统计和报告

import time
from locust import HttpUser, task, events
from locust.runners import STATE_STOPPING, STATE_STOPPED, STATE_CLEANUP

class AdvancedUser(HttpUser):
    @task
    def custom_task(self):
        start_time = time.time()
        try:
            response = self.client.get("/api/test")
            # 自定义成功条件
            if response.status_code == 200 and "success" in response.text:
                events.request.fire(
                    request_type="GET",
                    name="/api/test",
                    response_time=(time.time() - start_time) * 1000,
                    response_length=len(response.content),
                    exception=None,
                    context={}
                )
            else:
                events.request.fire(
                    request_type="GET",
                    name="/api/test",
                    response_time=(time.time() - start_time) * 1000,
                    response_length=0,
                    exception="Invalid response",
                    context={}
                )
        except Exception as e:
            events.request.fire(
                request_type="GET",
                name="/api/test",
                response_time=(time.time() - start_time) * 1000,
                response_length=0,
                exception=str(e),
                context={}
            )

七、运行 Locust

1. 命令行运行

# 基本运行（带 Web UI）
locust -f locustfile.py

# 无 Web UI 模式（用于 CI/CD）
locust -f locustfile.py --headless --users 100 --spawn-rate 10 --run-time 1h

# 指定主机
locust -f locustfile.py --host=http://localhost:8000

# 分布式运行
# 在主节点
locust -f locustfile.py --master

# 在工作节点
locust -f locustfile.py --worker --master-host=192.168.1.100

# 自定义端口
locust -f locustfile.py --web-port=8080

# 输出统计到 CSV
locust -f locustfile.py --csv=results --headless --users 100 --spawn-rate 10 --run-time 5m

2. 配置文件

创建 locust.conf 文件：

# locust.conf
locustfile = locustfile.py
host = http://localhost:8000
users = 100
spawn-rate = 10
run-time = 5m
headless = true
csv = results

运行：

locust --config=locust.conf

八、实战示例：API 负载测试

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

class ApiTestUser(HttpUser):
    wait_time = between(0.5, 2)
    host = "https://api.example.com"
    
    # 测试数据
    user_ids = list(range(1, 1000))
    product_ids = ["P1001", "P1002", "P1003", "P1004"]
    
    def on_start(self):
        """获取认证令牌"""
        response = self.client.post("/auth/token", json={
            "username": "test",
            "password": "test123"
        })
        self.token = response.json()["access_token"]
        self.headers = {"Authorization": f"Bearer {self.token}"}
    
    @tag("read")
    @task(5)
    def get_users(self):
        """获取用户列表"""
        params = {"limit": 10, "page": random.randint(1, 10)}
        self.client.get("/users", params=params, headers=self.headers)
    
    @tag("read")
    @task(3)
    def get_user_detail(self):
        """获取单个用户详情"""
        user_id = random.choice(self.user_ids)
        self.client.get(f"/users/{user_id}", headers=self.headers)
    
    @tag("write")
    @task(2)
    def create_order(self):
        """创建订单"""
        payload = {
            "product_id": random.choice(self.product_ids),
            "quantity": random.randint(1, 5),
            "user_id": random.choice(self.user_ids)
        }
        self.client.post("/orders", json=payload, headers=self.headers)
    
    @tag("write")
    @task(1)
    def update_user(self):
        """更新用户信息"""
        user_id = random.choice(self.user_ids)
        payload = {
            "email": f"user{user_id}@example.com",
            "name": f"User {user_id}"
        }
        self.client.put(f"/users/{user_id}", json=payload, headers=self.headers)
    
    @tag("performance")
    @task
    def complex_query(self):
        """复杂查询测试"""
        query = {
            "filters": {
                "status": ["active", "pending"],
                "date_range": {
                    "start": "2024-01-01",
                    "end": "2024-12-31"
                }
            },
            "sort": {"field": "created_at", "order": "desc"},
            "pagination": {"page": 1, "limit": 50}
        }
        self.client.post("/reports/query", json=query, headers=self.headers)

九、最佳实践

1. 测试设计原则

• 模拟真实用户行为：使用合理的等待时间和任务权重
• 参数化数据：避免缓存影响
• 渐进式加载：逐步增加用户数
• 设置断言：验证响应正确性

2. 性能监控要点

• 响应时间百分位：关注 p95, p99
• 错误率：控制在可接受范围
• 吞吐量：每秒请求数
• 资源使用率：CPU、内存、网络

3. 分布式测试配置

# docker-compose.yml 示例
version: '3'

services:
  master:
    image: locustio/locust
    ports:
      - "8089:8089"
    volumes:
      - ./locustfile:/mnt/locust
    command: -f /mnt/locust/locustfile.py --master -H http://target-service:8080

  worker:
    image: locustio/locust
    volumes:
      - ./locustfile:/mnt/locust
    command: -f /mnt/locust/locustfile.py --worker --master-host=master
    scale: 4  # 启动4个工作节点

十、与其他工具集成

1. 与 Docker 集成

# Dockerfile
FROM python:3.9-slim
RUN pip install locust
COPY locustfile.py /locustfile.py
WORKDIR /app
CMD ["locust", "-f", "locustfile.py", "--host=http://target:8080"]

2. 与 CI/CD 集成

# GitHub Actions 示例
name: Load Test
on: [push]

jobs:
  load-test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v2
      - name: Set up Python
        uses: actions/setup-python@v2
      - name: Install dependencies
        run: pip install locust
      - name: Run load test
        run: |
          locust -f locustfile.py \
            --headless \
            --users 100 \
            --spawn-rate 10 \
            --run-time 5m \
            --csv=results \
            --html=report.html
      - name: Upload results
        uses: actions/upload-artifact@v2
        with:
          name: load-test-results
          path: |
            results*.csv
            report.html