Locust脚本编写实战

Li JiaXin大约 5 分钟

Locust脚本编写实战

🎯 入门容易精通难！掌握了基础操作后，是时候学习如何编写真正实用的Locust脚本了。
就像从新手司机进阶到老司机，我们需要掌握各种复杂路况的驾驶技巧。
今天我来分享一些在实际项目中总结的脚本编写经验和技巧！💪

🏗️ User类深度解析

HttpUser类的核心属性

在Locust中，HttpUser就像是我们的"虚拟用户模板"，让我们深入了解它的各个组件：

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

class MyUser(HttpUser):
    # 等待时间策略
    wait_time = between(1, 3)  # 1-3秒随机等待
    
    # 目标主机
    host = "https://api.example.com"
    
    # 用户权重（多用户类时使用）
    weight = 3
    
    # 固定用户数（忽略权重）
    # fixed_count = 10

等待时间策略详解

等待时间就像真实用户的"思考时间"，不同策略适用于不同场景：

from locust import constant, between, constant_throughput, constant_pacing

class DifferentWaitStrategies(HttpUser):
    # 1. 固定等待时间
    wait_time = constant(2)  # 固定等待2秒
    
    # 2. 随机等待时间（最常用）
    wait_time = between(1, 5)  # 1-5秒随机
    
    # 3. 恒定吞吐量
    wait_time = constant_throughput(0.5)  # 每秒最多0.5个任务
    
    # 4. 恒定间隔
    wait_time = constant_pacing(3)  # 每3秒最多1个任务

Task任务系统

Task是Locust的核心，让我们看看各种定义方式：

class TaskExamples(HttpUser):
    wait_time = between(1, 2)
    host = "https://httpbin.org"
    
    # 方式1：装饰器方式（推荐）
    @task
    def simple_task(self):
        self.client.get("/get")
    
    # 方式2：带权重的任务
    @task(3)  # 权重为3，执行频率更高
    def weighted_task(self):
        self.client.get("/json")
    
    # 方式3：任务列表方式
    tasks = [simple_task, weighted_task]
    
    # 方式4：任务字典方式
    tasks = {simple_task: 1, weighted_task: 3}

🎭 生命周期钩子函数

on_start和on_stop

这两个函数就像用户的"登录"和"退出"操作：

class UserLifecycle(HttpUser):
    wait_time = between(1, 2)
    host = "https://api.example.com"
    
    def on_start(self):
        """用户启动时执行一次（类似登录）"""
        # 模拟用户登录
        response = self.client.post("/login", json={
            "username": "testuser",
            "password": "password123"
        })
        
        if response.status_code == 200:
            # 保存登录token
            self.token = response.json().get("token")
            print(f"用户登录成功，token: {self.token}")
        else:
            print("登录失败！")
    
    def on_stop(self):
        """用户停止时执行一次（类似退出）"""
        if hasattr(self, 'token'):
            self.client.post("/logout", headers={
                "Authorization": f"Bearer {self.token}"
            })
            print("用户已退出")
    
    @task
    def protected_api(self):
        """需要认证的API调用"""
        if hasattr(self, 'token'):
            self.client.get("/protected", headers={
                "Authorization": f"Bearer {self.token}"
            })

🔧 HTTP请求进阶技巧

请求参数和响应处理

class AdvancedRequests(HttpUser):
    wait_time = between(1, 2)
    host = "https://httpbin.org"
    
    @task
    def get_with_params(self):
        """GET请求带参数"""
        params = {
            "page": 1,
            "size": 10,
            "keyword": "locust"
        }
        response = self.client.get("/get", params=params)
        print(f"响应状态码: {response.status_code}")
        print(f"响应内容: {response.json()}")
    
    @task
    def post_with_json(self):
        """POST请求发送JSON数据"""
        data = {
            "name": "测试用户",
            "email": "test@example.com",
            "age": 25
        }
        response = self.client.post("/post", json=data)
        
        # 检查响应
        if response.status_code == 200:
            result = response.json()
            print(f"创建成功: {result}")
    
    @task
    def upload_file(self):
        """文件上传"""
        files = {
            'file': ('test.txt', 'Hello Locust!', 'text/plain')
        }
        response = self.client.post("/post", files=files)
        print(f"文件上传结果: {response.status_code}")

响应验证和错误处理

class ResponseValidation(HttpUser):
    wait_time = between(1, 2)
    host = "https://httpbin.org"
    
    @task
    def validate_response(self):
        """响应验证示例"""
        with self.client.get("/json", catch_response=True) as response:
            # 检查状态码
            if response.status_code != 200:
                response.failure(f"状态码错误: {response.status_code}")
                return
            
            # 检查响应时间
            if response.elapsed.total_seconds() > 2:
                response.failure("响应时间过长")
                return
            
            # 检查响应内容
            try:
                data = response.json()
                if "slideshow" not in data:
                    response.failure("响应内容格式错误")
                    return
                
                # 手动标记成功
                response.success()
                print("响应验证通过")
                
            except ValueError:
                response.failure("JSON解析失败")

📊 数据驱动测试

使用队列管理测试数据

from queue import Queue
import random

class DataDrivenTest(HttpUser):
    wait_time = between(1, 2)
    host = "https://httpbin.org"
    
    # 类级别的数据队列
    user_data_queue = Queue()
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 初始化测试数据
        if self.user_data_queue.empty():
            for i in range(100):
                self.user_data_queue.put({
                    "user_id": f"user_{i}",
                    "name": f"测试用户{i}",
                    "email": f"user{i}@test.com"
                })
    
    @task
    def test_with_unique_data(self):
        """使用唯一数据进行测试"""
        try:
            # 获取唯一数据
            user_data = self.user_data_queue.get(timeout=1)
            
            # 使用数据进行测试
            response = self.client.post("/post", json=user_data)
            
            if response.status_code == 200:
                print(f"用户 {user_data['user_id']} 测试成功")
            
            # 如果需要循环使用数据，可以重新放回队列
            # self.user_data_queue.put(user_data)
            
        except:
            print("数据队列已空，停止测试")
            self.environment.runner.quit()

随机数据生成

import random
import string
from faker import Faker

class RandomDataTest(HttpUser):
    wait_time = between(1, 2)
    host = "https://httpbin.org"
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.fake = Faker('zh_CN')  # 中文数据生成器
    
    def generate_random_user(self):
        """生成随机用户数据"""
        return {
            "name": self.fake.name(),
            "email": self.fake.email(),
            "phone": self.fake.phone_number(),
            "address": self.fake.address(),
            "company": self.fake.company()
        }
    
    @task
    def test_with_random_data(self):
        """使用随机数据测试"""
        user_data = self.generate_random_user()
        response = self.client.post("/post", json=user_data)
        
        if response.status_code == 200:
            print(f"随机用户 {user_data['name']} 创建成功")

🏷️ 标签系统

标签系统让我们可以选择性地运行特定的测试：

from locust import tag

class TaggedTests(HttpUser):
    wait_time = between(1, 2)
    host = "https://httpbin.org"
    
    @task
    @tag('smoke')
    def smoke_test(self):
        """冒烟测试"""
        self.client.get("/get")
    
    @task
    @tag('load')
    def load_test(self):
        """负载测试"""
        self.client.post("/post", json={"test": "data"})
    
    @task
    @tag('smoke', 'critical')
    def critical_test(self):
        """关键功能测试"""
        self.client.get("/status/200")

运行特定标签的测试：

# 只运行冒烟测试
locust --tags smoke

# 运行多个标签
locust --tags smoke load

# 排除特定标签
locust --exclude-tags load

🎯 实战技巧总结

1. 脚本组织建议

# 推荐的脚本结构
class BaseUser(HttpUser):
    """基础用户类，包含公共方法"""
    
    def login(self, username, password):
        """公共登录方法"""
        pass
    
    def logout(self):
        """公共退出方法"""
        pass

class NormalUser(BaseUser):
    """普通用户行为"""
    weight = 7
    
    @task(5)
    def browse_products(self):
        pass
    
    @task(2)
    def search_products(self):
        pass

class VIPUser(BaseUser):
    """VIP用户行为"""
    weight = 3
    
    @task(3)
    def browse_vip_products(self):
        pass

2. 性能优化建议

合理设置等待时间，模拟真实用户行为
使用连接池，避免频繁建立连接
适当使用缓存，减少重复计算
监控资源使用情况，避免客户端成为瓶颈

3. 调试技巧

# 开发调试时使用单用户模式
if __name__ == "__main__":
    import sys
    from locust.env import Environment
    from locust.stats import stats_printer
    
    # 创建环境
    env = Environment(user_classes=[MyUser])
    env.create_local_runner()
    
    # 启动单个用户进行调试
    env.runner.start(1, spawn_rate=1)
    
    # 运行一段时间后停止
    import time
    time.sleep(30)
    env.runner.quit()

🚀 下一步学习

掌握了脚本编写技巧后，建议继续学习：

Locust高级特性 - 事件系统、自定义统计等
分布式压测 - 多机器协同压测
性能调优 - 解决性能瓶颈问题

记住，好的压测脚本就像好的测试用例，要覆盖全面、逻辑清晰、易于维护！🎯