Locust框架封装与扩展
大约 8 分钟
Locust框架封装与扩展
🏗️ 用了一段时间Locust，发现每次都要写重复代码？是时候搭建自己的压测框架了！
就像从手工作坊升级到工业化生产，我们要让压测更加标准化、自动化。
今天分享如何基于Locust构建企业级压测框架，让你的团队效率翻倍！🚀
🎯 框架设计理念
设计目标
作为测试开发工程师，我们的框架要解决这些痛点：
- 降低门槛：让不熟悉Locust的同事也能快速上手
- 提高复用：公共功能封装，避免重复造轮子
- 统一标准：规范化测试流程和报告格式
- 易于扩展：支持各种业务场景的定制需求
架构设计
📦 LocustFramework
├── 🏗️ core/ # 核心框架
│ ├── base_user.py # 基础用户类
│ ├── auth_manager.py # 认证管理
│ ├── data_manager.py # 数据管理
│ └── report_manager.py # 报告管理
├── 🔧 utils/ # 工具模块
│ ├── config.py # 配置管理
│ ├── logger.py # 日志管理
│ └── helpers.py # 辅助函数
├── 📊 plugins/ # 插件系统
│ ├── monitoring.py # 监控插件
│ ├── alerting.py # 告警插件
│ └── custom_stats.py # 自定义统计
├── 🎨 templates/ # 测试模板
│ ├── api_test.py # API测试模板
│ ├── web_test.py # Web测试模板
│ └── load_shapes.py # 负载模型模板
└── 📋 examples/ # 示例项目
├── simple_api/ # 简单API测试
├── complex_workflow/ # 复杂业务流程
└── microservice/ # 微服务测试🏗️ 核心框架实现
1. 基础用户类封装
# core/base_user.py
from locust import HttpUser, task, events
from abc import ABC, abstractmethod
import time
import json
from utils.logger import get_logger
from utils.config import Config
class BaseUser(HttpUser, ABC):
"""基础用户类，提供通用功能"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.logger = get_logger(self.__class__.__name__)
self.config = Config()
self.user_data = {}
self.session_data = {}
def on_start(self):
"""用户启动时的通用初始化"""
self.logger.info(f"🚀 用户 {self.get_user_id()} 开始测试")
# 设置通用请求头
self.client.headers.update({
'User-Agent': f'LocustFramework/{self.config.version}',
'Accept': 'application/json',
'Content-Type': 'application/json'
})
# 执行自定义初始化
self.setup()
def on_stop(self):
"""用户停止时的通用清理"""
self.logger.info(f"🛑 用户 {self.get_user_id()} 结束测试")
self.cleanup()
@abstractmethod
def setup(self):
"""子类实现的初始化方法"""
pass
@abstractmethod
def cleanup(self):
"""子类实现的清理方法"""
pass
def get_user_id(self):
"""获取用户唯一标识"""
return f"{self.__class__.__name__}_{id(self)}"
def safe_request(self, method, url, name=None, **kwargs):
"""安全的请求方法，包含错误处理和重试"""
max_retries = self.config.max_retries
retry_delay = self.config.retry_delay
for attempt in range(max_retries + 1):
try:
start_time = time.time()
with self.client.request(
method, url, name=name, catch_response=True, **kwargs
) as response:
# 记录请求详情
self._log_request(method, url, response, start_time)
# 验证响应
if self._validate_response(response):
response.success()
return response
else:
response.failure("Response validation failed")
except Exception as e:
self.logger.error(f"❌ 请求失败 (尝试 {attempt + 1}/{max_retries + 1}): {e}")
if attempt < max_retries:
time.sleep(retry_delay)
continue
else:
raise
return None
def _log_request(self, method, url, response, start_time):
"""记录请求日志"""
duration = (time.time() - start_time) * 1000
self.logger.debug(
f"📡 {method} {url} -> {response.status_code} ({duration:.1f}ms)"
)
def _validate_response(self, response):
"""验证响应（子类可重写）"""
return 200 <= response.status_code < 400
class APIUser(BaseUser):
"""API测试用户基类"""
def setup(self):
"""API用户初始化"""
# 执行认证
if hasattr(self, 'authenticate'):
self.authenticate()
def cleanup(self):
"""API用户清理"""
# 执行登出
if hasattr(self, 'logout'):
self.logout()
def get_json_response(self, response):
"""安全获取JSON响应"""
try:
return response.json()
except json.JSONDecodeError:
self.logger.error("❌ JSON解析失败")
return None
class WebUser(BaseUser):
"""Web测试用户基类"""
def setup(self):
"""Web用户初始化"""
# 设置Web特定的请求头
self.client.headers.update({
'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
'Accept-Encoding': 'gzip, deflate'
})
def cleanup(self):
"""Web用户清理"""
pass
def extract_csrf_token(self, response):
"""提取CSRF令牌"""
from bs4 import BeautifulSoup
soup = BeautifulSoup(response.text, 'html.parser')
csrf_input = soup.find('input', {'name': 'csrf_token'})
return csrf_input['value'] if csrf_input else None2. 认证管理器
# core/auth_manager.py
import time
import hmac
import hashlib
import base64
from requests.auth import AuthBase
from utils.logger import get_logger
class AuthManager:
"""认证管理器"""
def __init__(self):
self.logger = get_logger(self.__class__.__name__)
self.auth_cache = {}
def get_auth(self, auth_type, **kwargs):
"""获取认证对象"""
auth_map = {
'basic': BasicAuth,
'bearer': BearerAuth,
'hmac': HMACAuth,
'oauth2': OAuth2Auth
}
auth_class = auth_map.get(auth_type)
if not auth_class:
raise ValueError(f"不支持的认证类型: {auth_type}")
return auth_class(**kwargs)
class BasicAuth(AuthBase):
"""基础认证"""
def __init__(self, username, password):
self.username = username
self.password = password
def __call__(self, request):
credentials = f"{self.username}:{self.password}"
encoded = base64.b64encode(credentials.encode()).decode()
request.headers['Authorization'] = f'Basic {encoded}'
return request
class BearerAuth(AuthBase):
"""Bearer Token认证"""
def __init__(self, token):
self.token = token
def __call__(self, request):
request.headers['Authorization'] = f'Bearer {self.token}'
return request
class HMACAuth(AuthBase):
"""HMAC签名认证"""
def __init__(self, access_key, secret_key):
self.access_key = access_key
self.secret_key = secret_key
def __call__(self, request):
timestamp = str(int(time.time()))
# 构造签名字符串
string_to_sign = f"{request.method}\n{request.url}\n{timestamp}"
# 生成签名
signature = hmac.new(
self.secret_key.encode(),
string_to_sign.encode(),
hashlib.sha256
).hexdigest()
# 添加认证头
request.headers.update({
'X-Access-Key': self.access_key,
'X-Timestamp': timestamp,
'X-Signature': signature
})
return request
class OAuth2Auth(AuthBase):
"""OAuth2认证"""
def __init__(self, client_id, client_secret, token_url):
self.client_id = client_id
self.client_secret = client_secret
self.token_url = token_url
self.access_token = None
self.token_expires = 0
def __call__(self, request):
if self._token_expired():
self._refresh_token()
if self.access_token:
request.headers['Authorization'] = f'Bearer {self.access_token}'
return request
def _token_expired(self):
"""检查token是否过期"""
return time.time() >= self.token_expires
def _refresh_token(self):
"""刷新访问令牌"""
import requests
data = {
'grant_type': 'client_credentials',
'client_id': self.client_id,
'client_secret': self.client_secret
}
response = requests.post(self.token_url, data=data)
if response.status_code == 200:
token_data = response.json()
self.access_token = token_data['access_token']
expires_in = token_data.get('expires_in', 3600)
self.token_expires = time.time() + expires_in - 60 # 提前60秒刷新3. 数据管理器
# core/data_manager.py
import csv
import json
import random
import threading
from queue import Queue, Empty
from faker import Faker
from utils.logger import get_logger
class DataManager:
"""数据管理器"""
def __init__(self):
self.logger = get_logger(self.__class__.__name__)
self.data_sources = {}
self.lock = threading.Lock()
def register_data_source(self, name, source):
"""注册数据源"""
with self.lock:
self.data_sources[name] = source
self.logger.info(f"📊 注册数据源: {name}")
def get_data(self, source_name, **kwargs):
"""获取数据"""
source = self.data_sources.get(source_name)
if not source:
raise ValueError(f"数据源不存在: {source_name}")
return source.get_data(**kwargs)
class CSVDataSource:
"""CSV数据源"""
def __init__(self, file_path, unique=False):
self.file_path = file_path
self.unique = unique
self.data_queue = Queue()
self.original_data = []
self._load_data()
def _load_data(self):
"""加载CSV数据"""
with open(self.file_path, 'r', encoding='utf-8') as f:
reader = csv.DictReader(f)
for row in reader:
self.original_data.append(row)
self.data_queue.put(row)
print(f"📁 加载CSV数据: {len(self.original_data)} 条记录")
def get_data(self, **kwargs):
"""获取数据"""
try:
data = self.data_queue.get(timeout=1)
# 如果不要求唯一性，使用后放回队列
if not self.unique:
self.data_queue.put(data)
return data
except Empty:
if self.unique:
raise ValueError("数据已用完")
else:
# 重新加载数据
for row in self.original_data:
self.data_queue.put(row)
return self.data_queue.get()
class RandomDataSource:
"""随机数据源"""
def __init__(self, locale='zh_CN'):
self.fake = Faker(locale)
def get_data(self, data_type='user', **kwargs):
"""生成随机数据"""
generators = {
'user': self._generate_user,
'product': self._generate_product,
'order': self._generate_order
}
generator = generators.get(data_type)
if not generator:
raise ValueError(f"不支持的数据类型: {data_type}")
return generator(**kwargs)
def _generate_user(self, **kwargs):
"""生成用户数据"""
return {
'id': self.fake.uuid4(),
'name': self.fake.name(),
'email': self.fake.email(),
'phone': self.fake.phone_number(),
'address': self.fake.address(),
'company': self.fake.company(),
'job': self.fake.job()
}
def _generate_product(self, **kwargs):
"""生成产品数据"""
return {
'id': self.fake.uuid4(),
'name': self.fake.catch_phrase(),
'description': self.fake.text(max_nb_chars=200),
'price': round(random.uniform(10, 1000), 2),
'category': random.choice(['电子', '服装', '食品', '图书', '家居']),
'stock': random.randint(0, 1000)
}
def _generate_order(self, **kwargs):
"""生成订单数据"""
return {
'id': self.fake.uuid4(),
'user_id': kwargs.get('user_id', self.fake.uuid4()),
'product_ids': [self.fake.uuid4() for _ in range(random.randint(1, 5))],
'total_amount': round(random.uniform(50, 2000), 2),
'status': random.choice(['pending', 'paid', 'shipped', 'delivered']),
'created_at': self.fake.date_time_this_year().isoformat()
}
class DatabaseDataSource:
"""数据库数据源"""
def __init__(self, connection_string, query):
self.connection_string = connection_string
self.query = query
self.connection = None
self._connect()
def _connect(self):
"""连接数据库"""
# 这里可以根据需要支持不同的数据库
import sqlite3
self.connection = sqlite3.connect(self.connection_string)
self.connection.row_factory = sqlite3.Row
def get_data(self, **kwargs):
"""从数据库获取数据"""
cursor = self.connection.cursor()
cursor.execute(self.query, kwargs)
row = cursor.fetchone()
return dict(row) if row else None🔧 工具模块实现
配置管理
# utils/config.py
import os
import json
import yaml
from pathlib import Path
class Config:
"""配置管理类"""
def __init__(self, config_file=None):
self.config_file = config_file or self._find_config_file()
self.config_data = {}
self._load_config()
def _find_config_file(self):
"""查找配置文件"""
possible_files = [
'config.yaml',
'config.yml',
'config.json',
'locust.conf'
]
for file_name in possible_files:
if Path(file_name).exists():
return file_name
return None
def _load_config(self):
"""加载配置"""
if not self.config_file:
self._load_default_config()
return
file_ext = Path(self.config_file).suffix.lower()
with open(self.config_file, 'r', encoding='utf-8') as f:
if file_ext in ['.yaml', '.yml']:
self.config_data = yaml.safe_load(f)
elif file_ext == '.json':
self.config_data = json.load(f)
else:
self._load_default_config()
def _load_default_config(self):
"""加载默认配置"""
self.config_data = {
'version': '1.0.0',
'max_retries': 3,
'retry_delay': 1,
'timeout': 30,
'log_level': 'INFO',
'report_format': 'html'
}
def get(self, key, default=None):
"""获取配置值"""
keys = key.split('.')
value = self.config_data
for k in keys:
if isinstance(value, dict) and k in value:
value = value[k]
else:
return default
return value
def __getattr__(self, name):
"""支持属性访问"""
return self.get(name)日志管理
# utils/logger.py
import logging
import sys
from pathlib import Path
from datetime import datetime
def get_logger(name, level=logging.INFO):
"""获取日志记录器"""
logger = logging.getLogger(name)
if logger.handlers:
return logger
logger.setLevel(level)
# 创建格式化器
formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
# 控制台处理器
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setFormatter(formatter)
logger.addHandler(console_handler)
# 文件处理器
log_dir = Path('logs')
log_dir.mkdir(exist_ok=True)
file_handler = logging.FileHandler(
log_dir / f'locust_{datetime.now().strftime("%Y%m%d")}.log',
encoding='utf-8'
)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
return logger🎨 测试模板
API测试模板
# templates/api_test.py
from core.base_user import APIUser
from core.auth_manager import AuthManager
from core.data_manager import DataManager, RandomDataSource
from locust import task, between
class APITestTemplate(APIUser):
"""API测试模板"""
wait_time = between(1, 3)
host = "https://api.example.com"
def setup(self):
"""初始化设置"""
super().setup()
# 设置认证
auth_manager = AuthManager()
self.client.auth = auth_manager.get_auth(
'bearer',
token='your_api_token'
)
# 设置数据源
self.data_manager = DataManager()
self.data_manager.register_data_source(
'random_user',
RandomDataSource()
)
def cleanup(self):
"""清理资源"""
super().cleanup()
@task(3)
def get_users(self):
"""获取用户列表"""
response = self.safe_request('GET', '/api/users')
if response:
users = self.get_json_response(response)
self.logger.info(f"📋 获取到 {len(users)} 个用户")
@task(2)
def create_user(self):
"""创建用户"""
user_data = self.data_manager.get_data('random_user', data_type='user')
response = self.safe_request(
'POST',
'/api/users',
json=user_data,
name='create_user'
)
if response:
created_user = self.get_json_response(response)
self.logger.info(f"✅ 创建用户成功: {created_user.get('id')}")
@task(1)
def update_user(self):
"""更新用户"""
user_id = "test_user_id" # 实际项目中应该从数据源获取
update_data = {
"name": "Updated Name",
"email": "updated@example.com"
}
response = self.safe_request(
'PUT',
f'/api/users/{user_id}',
json=update_data,
name='update_user'
)
if response:
self.logger.info(f"✅ 更新用户成功: {user_id}")🚀 使用示例
快速开始
# examples/quick_start.py
from templates.api_test import APITestTemplate
from locust import between
class MyAPITest(APITestTemplate):
"""我的API测试"""
wait_time = between(0.5, 2)
host = "https://my-api.com"
def setup(self):
"""自定义初始化"""
super().setup()
# 添加自定义初始化逻辑
@task
def custom_endpoint(self):
"""自定义端点测试"""
self.safe_request('GET', '/api/custom')
# 运行命令：locust -f examples/quick_start.py🎯 框架优势总结
通过这个框架封装，我们实现了：
- 标准化：统一的代码结构和规范
- 复用性：公共功能可重复使用
- 扩展性：支持插件和自定义扩展
- 易用性：降低了使用门槛
- 维护性：代码结构清晰，易于维护
这个框架就像搭积木一样，你可以根据需要组合不同的组件，快速构建出适合自己业务的压测方案！🧱
下一步我们将看到这个框架在实际项目中的应用案例。