Locust Python 压测实战

# 创建虚拟环境
python -m venv .venv
source .venv/bin/activate  # Windows: .venv\Scripts\activate

# 安装 Locust（含 Web UI 与分布式支持）
pip install locust==2.44.1

# 验证版本
locust --version
# locust 2.44.1

# locustfile.py
from locust import HttpUser, task, between

class WebsiteUser(HttpUser):
    """模拟一个访问网站的用户"""

    # 两次请求之间的思考时间：1~3 秒随机
    wait_time = between(1, 3)

    # host 指定被测系统的根地址，也可通过 --host 命令行参数覆盖
    host = "http://localhost:8080"

    @task(2)  # 权重为 2，被调用的频率更高
    def list_courses(self):
        # self.client 是 requests 风格的客户端，自动记录响应时间与状态
        self.client.get("/api/courses")

    @task(1)  # 权重为 1
    def login(self):
        self.client.post("/api/login", json={
            "username": "tester",
            "password": "pass1234",
        })

locust -f locustfile.py
# 默认监听 http://localhost:8089

from locust import HttpUser, task, between

class ApiUser(HttpUser):
    wait_time = between(1, 2)

    @task
    def get_profile(self):
        # GET 请求，自动记录响应时间
        self.client.get("/api/profile")

    @task
    def create_order(self):
        # POST 请求，json 参数会自动设置 Content-Type
        self.client.post("/api/orders", json={"sku_id": 1001, "qty": 2})

    @task
    def upload_file(self):
        # 上传文件：与 requests 用法一致
        with open("avatar.png", "rb") as f:
            self.client.post("/api/upload", files={"file": f})

from locust import HttpUser, TaskSet, task, between

class OrderFlow(TaskSet):
    """下单完整链路：登录 → 浏览商品 → 加购 → 下单"""

    @task
    def login(self):
        resp = self.client.post("/api/login", json={
            "username": "buyer",
            "password": "pass1234",
        })
        # 提取 token 后续接口使用（关联）
        self.token = resp.json()["token"]

    @task(3)
    def browse_product(self):
        self.client.get("/api/products/1001",
                        headers={"Authorization": f"Bearer {self.token}"})

    @task(2)
    def add_to_cart(self):
        self.client.post("/api/cart", json={"sku_id": 1001, "qty": 1},
                         headers={"Authorization": f"Bearer {self.token}"})

    @task(1)
    def place_order(self):
        self.client.post("/api/orders",
                         headers={"Authorization": f"Bearer {self.token}"})
        # 5% 概率退出当前 TaskSet，回到父级
        if randint(0, 99) < 5:
            self.interrupt()

class WebsiteUser(HttpUser):
    tasks = [OrderFlow]
    wait_time = between(1, 3)

from locust import HttpUser, between, constant, constant_pacing

class UserA(HttpUser):
    wait_time = between(1, 5)        # 每次任务后随机等待 1~5 秒

class UserB(HttpUser):
    wait_time = constant(2)          # 每次任务后固定等待 2 秒

class UserC(HttpUser):
    wait_time = constant_pacing(3)   # 任务总间隔（含执行时间）保持 3 秒

import csv
import itertools
from locust import HttpUser, task, between, events

# 全局数据池，所有用户共享
user_credentials = []

@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    """测试开始前加载测试数据到内存"""
    global user_credentials
    with open("users.csv", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        user_credentials = list(reader)

class ApiUser(HttpUser):
    wait_time = between(1, 2)
    credentials_iter = None  # 类变量，作为共享迭代器

    def on_start(self):
        """每个用户启动时分配一组凭据"""
        # 通过类级别的迭代器实现"每个用户不同数据"
        if ApiUser.credentials_iter is None:
            ApiUser.credentials_iter = itertools.cycle(user_credentials)
        self.cred = next(ApiUser.credentials_iter)

    @task
    def login(self):
        self.client.post("/api/login", json={
            "username": self.cred["username"],
            "password": self.cred["password"],
        })

import re
from locust import HttpUser, task, between

class OrderUser(HttpUser):
    wait_time = between(1, 2)

    def on_start(self):
        # 登录获取 token
        resp = self.client.post("/api/login", json={
            "username": "tester", "password": "pass1234"})
        self.token = resp.json()["data"]["token"]

    @task
    def create_and_pay(self):
        # 创建订单，提取 order_id
        resp = self.client.post("/api/orders",
                                json={"sku_id": 1001},
                                headers={"Authorization": f"Bearer {self.token}"})
        order_id = resp.json()["data"]["order_id"]

        # 使用 order_id 调用支付接口
        self.client.post(f"/api/orders/{order_id}/pay",
                         headers={"Authorization": f"Bearer {self.token}"})

    @task
    def search_with_regex(self):
        # 正则提取示例：从 HTML 响应中提取 csrf_token
        resp = self.client.get("/form")
        match = re.search(r'name="csrf_token" value="([^"]+)"', resp.text)
        if match:
            csrf = match.group(1)
            self.client.post("/submit", data={"csrf_token": csrf})

import time
import grpc
from locust import User, task, between, events

class GrpcClient:
    """封装 grpc.Channel，使其在调用时触发 Locust 事件"""

    def __init__(self, environment, stub_cls):
        self.env = environment
        self.channel = grpc.insecure_channel("localhost:50051")
        self.stub = stub_cls(self.channel)

    def send(self, name, request, timeout=5):
        """统一发送方法，记录请求耗时与结果"""
        start = time.time()
        try:
            response = self.stub.GetFeature(request, timeout=timeout)
            elapsed = int((time.time() - start) * 1000)
            # 触发请求成功事件，被 Locust 统计系统捕获
            events.request.fire(
                request_type="gRPC", name=name,
                response_time=elapsed, response_length=len(response.SerializeToString()),
                exception=None, context=self.env,
            )
            return response
        except grpc.RpcError as e:
            elapsed = int((time.time() - start) * 1000)
            events.request.fire(
                request_type="gRPC", name=name,
                response_time=elapsed, response_length=0,
                exception=e, context=self.env,
            )
            raise

# 假设已通过 grpcio-tools 生成 route_guide_pb2 / route_guide_pb2_grpc
import route_guide_pb2_grpc
import route_guide_pb2

class GrpcUser(User):
    abstract = True  # 抽象基类，不会被实例化为虚拟用户
    stub_cls = None

    def __init__(self, environment):
        super().__init__(environment)
        self.client = GrpcClient(environment, self.stub_cls)

class RouteGuideUser(GrpcUser):
    wait_time = between(1, 2)
    stub_cls = route_guide_pb2_grpc.RouteGuideStub

    @task
    def get_feature(self):
        point = route_guide_pb2.Point(latitude=409146138, longitude=-746188906)
        self.client.send("GetFeature", point)

import logging
from locust import events

# 请求级：每个请求完成后触发
@events.request.add_listener
def on_request(request_type, name, response_time, response_length, exception, **kwargs):
    if exception:
        logging.warning("请求失败 %s %s: %s", request_type, name, exception)
    elif response_time > 2000:
        logging.info("慢请求 %s %s: %dms", request_type, name, response_time)

# 测试级：压测开始前触发（每个 Worker 各执行一次）
@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    logging.info("压测启动，目标用户数: %d", environment.parsed_options.num_users)

# 测试级：压测结束后触发
@events.test_stop.add_listener
def on_test_stop(environment, **kwargs):
    logging.info("压测结束")

# 用户级：每个用户启动时触发
@events.user_class_start.add_listener
def on_user_class_start(user_class, **kwargs):
    logging.info("用户类 %s 开始实例化", user_class.__name__)

locust -f locustfile.py --master --expect-workers 3
# Master 监听 8089 端口等待 Worker 连接

locust -f locustfile.py --worker --master-host=192.168.1.10
# Worker 连接到 Master 后等待任务分发

# docker-compose.yml
version: "3.8"

services:
  master:
    image: locustio/locust:2.44.1
    ports:
      - "8089:8089"
    volumes:
      - ./locustfile.py:/mnt/locust/locustfile.py
    command: >
      -f /mnt/locust/locustfile.py
      --master
      --expect-workers 3
      --host http://target-service:8080

  worker:
    image: locustio/locust:2.44.1
    volumes:
      - ./locustfile.py:/mnt/locust/locustfile.py
    deploy:
      replicas: 3
    command: >
      -f /mnt/locust/locustfile.py
      --worker
      --master-host master
    depends_on:
      - master

docker compose up -d
# 浏览器访问 http://localhost:8089 开始压测

locust -f locustfile.py \
  --headless \
  --users 500 \              # 总用户数
  --spawn-rate 50 \          # 每秒启动用户数
  --run-time 5m \            # 持续 5 分钟
  --host http://staging.example.com \
  --html report.html \       # 输出 HTML 报告
  --csv result               # 输出 CSV 结果（result_stats.csv 等）

# .github/workflows/load-test.yml
name: Load Test

on:
  push:
    branches: [ main ]
  workflow_dispatch:

jobs:
  locust:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4

      - name: Setup Python
        uses: actions/setup-python@v5
        with:
          python-version: "3.11"

      - name: Install Locust
        run: pip install locust==2.44.1

      - name: Run Load Test
        run: |
          locust -f tests/load/locustfile.py \
            --headless \
            --users 200 \
            --spawn-rate 20 \
            --run-time 3m \
            --host https://staging.example.com \
            --csv report \
            --html report.html \
            --fail-on-error

      - name: Upload Report
        if: always()
        uses: actions/upload-artifact@v4
        with:
          name: locust-report
          path: |
            report.html
            report_stats.csv

from locust import LoadTestShape

class StepLoadShape(LoadTestShape):
    """阶梯加压：每 60 秒增加 100 用户，达到 500 后保持 5 分钟"""

    stages = [
        {"duration": 60,  "users": 100, "spawn_rate": 10},
        {"duration": 120, "users": 200, "spawn_rate": 10},
        {"duration": 180, "users": 300, "spawn_rate": 10},
        {"duration": 240, "users": 400, "spawn_rate": 10},
        {"duration": 540, "users": 500, "spawn_rate": 10},  # 持续 5 分钟
    ]

    def tick(self):
        run_time = self.get_run_time()
        for stage in self.stages:
            if run_time < stage["duration"]:
                return (stage["users"], stage["spawn_rate"])
        return None  # 返回 None 表示压测结束