压测脚本引擎设计与实现
大约 10 分钟
压测脚本引擎设计与实现
前言：脚本引擎的"魔法"
还记得我刚开始写压测脚本时，每个接口都要写一遍重复的代码，就像每次做饭都要从洗菜开始一样繁琐。后来我意识到，好的脚本引擎就像一个智能的厨师助手，把复杂的准备工作都做好了，让你专注于创造美味的测试场景。
今天我们就来深入探讨如何设计和实现一个灵活而强大的压测脚本引擎，让复杂的业务场景测试变得简单而优雅。这不是简单的代码复制粘贴，而是基于真实项目经验的深度设计实践。
脚本引擎设计理念：简单而强大
设计目标
"""
脚本引擎的设计目标
就像设计一辆好车，既要操作简单，又要性能强劲
"""
design_goals = {
"易用性": "简单的API，复杂的功能",
"可复用性": "一次编写，到处使用",
"可扩展性": "插件机制，自定义扩展",
"可维护性": "清晰的结构，易于调试",
"高性能": "优化的执行，资源高效"
}核心特性一览
# core/script_engine.py 核心特性
class ScriptEngine:
"""
压测脚本引擎 - 业务场景的"导演"
特性清单：
✅ HTTP客户端封装
✅ 请求/响应处理
✅ 断言验证机制
✅ 数据提取与传递
✅ 错误处理与重试
✅ 性能监控
✅ 业务流程编排
✅ 插件扩展支持
"""
def __init__(self, user_instance):
self.user = user_instance
self.client = user_instance.client
self.session_data = {}
self.performance_data = []
# 初始化各个组件
self._setup_http_client()
self._setup_validators()
self._setup_extractors()
self._setup_monitors()
logger.info(f"🚀 脚本引擎初始化完成")HTTP客户端封装：网络请求的"瑞士军刀"
1. 增强的HTTP客户端
# core/http_client.py
"""
增强的HTTP客户端 - 网络请求的专家
让HTTP请求变得简单而强大
"""
import time
import json
from typing import Dict, Any, Optional, Union
from urllib.parse import urljoin, urlparse
import requests
from locust.clients import HttpSession
class EnhancedHttpClient:
"""
增强的HTTP客户端
在Locust原生客户端基础上，添加更多实用功能
"""
def __init__(self, base_client: HttpSession):
self.base_client = base_client
self.session_data = {}
self.default_headers = {}
self.request_hooks = []
self.response_hooks = []
def set_default_headers(self, headers: Dict[str, str]):
"""设置默认请求头"""
self.default_headers.update(headers)
def add_request_hook(self, hook_func):
"""添加请求前置钩子"""
self.request_hooks.append(hook_func)
def add_response_hook(self, hook_func):
"""添加响应后置钩子"""
self.response_hooks.append(hook_func)
def request(self, method: str, url: str, **kwargs) -> 'EnhancedResponse':
"""
发送HTTP请求
Args:
method: HTTP方法
url: 请求URL
**kwargs: 其他请求参数
Returns:
增强的响应对象
"""
# 合并默认头部
headers = self.default_headers.copy()
headers.update(kwargs.get('headers', {}))
kwargs['headers'] = headers
# 执行请求前置钩子
for hook in self.request_hooks:
kwargs = hook(method, url, kwargs) or kwargs
# 记录请求开始时间
start_time = time.time()
# 发送请求
response = self.base_client.request(method, url, **kwargs)
# 计算响应时间
response_time = (time.time() - start_time) * 1000
# 创建增强响应对象
enhanced_response = EnhancedResponse(response, response_time)
# 执行响应后置钩子
for hook in self.response_hooks:
enhanced_response = hook(enhanced_response) or enhanced_response
return enhanced_response
def get(self, url: str, **kwargs) -> 'EnhancedResponse':
"""GET请求"""
return self.request('GET', url, **kwargs)
def post(self, url: str, **kwargs) -> 'EnhancedResponse':
"""POST请求"""
return self.request('POST', url, **kwargs)
def put(self, url: str, **kwargs) -> 'EnhancedResponse':
"""PUT请求"""
return self.request('PUT', url, **kwargs)
def delete(self, url: str, **kwargs) -> 'EnhancedResponse':
"""DELETE请求"""
return self.request('DELETE', url, **kwargs)
def patch(self, url: str, **kwargs) -> 'EnhancedResponse':
"""PATCH请求"""
return self.request('PATCH', url, **kwargs)
class EnhancedResponse:
"""
增强的响应对象
在原生响应基础上，添加更多便利方法
"""
def __init__(self, response: requests.Response, response_time: float):
self.response = response
self.response_time = response_time
self._json_data = None
@property
def status_code(self) -> int:
"""状态码"""
return self.response.status_code
@property
def headers(self) -> Dict[str, str]:
"""响应头"""
return dict(self.response.headers)
@property
def text(self) -> str:
"""响应文本"""
return self.response.text
@property
def content(self) -> bytes:
"""响应内容"""
return self.response.content
def json(self) -> Dict[str, Any]:
"""解析JSON响应"""
if self._json_data is None:
try:
self._json_data = self.response.json()
except ValueError:
raise ValueError("响应不是有效的JSON格式")
return self._json_data
def is_success(self) -> bool:
"""判断请求是否成功"""
return 200 <= self.status_code < 300
def is_client_error(self) -> bool:
"""判断是否为客户端错误"""
return 400 <= self.status_code < 500
def is_server_error(self) -> bool:
"""判断是否为服务器错误"""
return 500 <= self.status_code < 600
def extract_data(self, path: str, default=None):
"""使用JSONPath提取数据"""
try:
import jsonpath_ng
jsonpath_expr = jsonpath_ng.parse(path)
matches = jsonpath_expr.find(self.json())
return matches[0].value if matches else default
except Exception:
return default
def assert_status_code(self, expected_code: int):
"""断言状态码"""
if self.status_code != expected_code:
raise AssertionError(f"期望状态码 {expected_code}，实际 {self.status_code}")
return self
def assert_response_time(self, max_time: float):
"""断言响应时间"""
if self.response_time > max_time:
raise AssertionError(f"响应时间 {self.response_time}ms 超过限制 {max_time}ms")
return self
def assert_json_path(self, path: str, expected_value):
"""断言JSON路径值"""
actual_value = self.extract_data(path)
if actual_value != expected_value:
raise AssertionError(f"路径 {path} 期望值 {expected_value}，实际值 {actual_value}")
return self2. 业务流程建模
# core/business_flow.py
"""
业务流程建模 - 复杂场景的"编剧"
将复杂的业务流程抽象为可复用的组件
"""
from typing import List, Dict, Any, Callable
from abc import ABC, abstractmethod
class FlowStep(ABC):
"""
流程步骤基类
每个步骤都是业务流程中的一个环节
"""
def __init__(self, name: str, description: str = ""):
self.name = name
self.description = description
self.retry_count = 0
self.max_retries = 3
@abstractmethod
def execute(self, context: Dict[str, Any]) -> Dict[str, Any]:
"""执行步骤"""
pass
def on_success(self, context: Dict[str, Any], result: Any):
"""成功回调"""
pass
def on_failure(self, context: Dict[str, Any], error: Exception):
"""失败回调"""
pass
def should_retry(self, error: Exception) -> bool:
"""判断是否应该重试"""
return self.retry_count < self.max_retries
class HttpRequestStep(FlowStep):
"""
HTTP请求步骤
封装HTTP请求的通用逻辑
"""
def __init__(self, name: str, method: str, url: str, **kwargs):
super().__init__(name)
self.method = method
self.url = url
self.request_kwargs = kwargs
self.validators = []
self.extractors = []
def add_validator(self, validator_func: Callable):
"""添加响应验证器"""
self.validators.append(validator_func)
return self
def add_extractor(self, key: str, path: str):
"""添加数据提取器"""
self.extractors.append((key, path))
return self
def execute(self, context: Dict[str, Any]) -> Dict[str, Any]:
"""执行HTTP请求"""
client = context['client']
# 处理URL模板
url = self.url.format(**context)
# 处理请求参数模板
request_kwargs = self._process_templates(self.request_kwargs, context)
# 发送请求
response = client.request(self.method, url, **request_kwargs)
# 验证响应
for validator in self.validators:
validator(response)
# 提取数据
extracted_data = {}
for key, path in self.extractors:
extracted_data[key] = response.extract_data(path)
# 更新上下文
context.update(extracted_data)
context['last_response'] = response
return context
def _process_templates(self, data: Any, context: Dict[str, Any]) -> Any:
"""处理模板变量"""
if isinstance(data, str):
return data.format(**context)
elif isinstance(data, dict):
return {k: self._process_templates(v, context) for k, v in data.items()}
elif isinstance(data, list):
return [self._process_templates(item, context) for item in data]
else:
return data
class BusinessFlow:
"""
业务流程编排器
将多个步骤组织成完整的业务流程
"""
def __init__(self, name: str):
self.name = name
self.steps: List[FlowStep] = []
self.context = {}
def add_step(self, step: FlowStep):
"""添加流程步骤"""
self.steps.append(step)
return self
def set_context(self, **kwargs):
"""设置流程上下文"""
self.context.update(kwargs)
return self
def execute(self, client) -> Dict[str, Any]:
"""执行完整流程"""
self.context['client'] = client
for step in self.steps:
try:
self.context = step.execute(self.context)
step.on_success(self.context, self.context.get('last_response'))
except Exception as e:
step.on_failure(self.context, e)
if step.should_retry(e):
step.retry_count += 1
# 重试逻辑
continue
else:
raise
return self.context
# 业务流程示例
class ECommerceFlow:
"""电商业务流程示例"""
@staticmethod
def create_user_registration_flow() -> BusinessFlow:
"""用户注册流程"""
flow = BusinessFlow("用户注册流程")
# 步骤1：获取验证码
get_captcha = HttpRequestStep(
"获取验证码",
"GET",
"/api/captcha"
).add_extractor("captcha_token", "$.data.token")
# 步骤2：注册用户
register_user = HttpRequestStep(
"注册用户",
"POST",
"/api/register",
json={
"username": "{username}",
"password": "{password}",
"email": "{email}",
"captcha_token": "{captcha_token}"
}
).add_validator(lambda r: r.assert_status_code(201)) \
.add_extractor("user_id", "$.data.user_id")
# 步骤3：发送欢迎邮件
send_welcome = HttpRequestStep(
"发送欢迎邮件",
"POST",
"/api/emails/welcome",
json={"user_id": "{user_id}"}
).add_validator(lambda r: r.assert_status_code(200))
return flow.add_step(get_captcha) \
.add_step(register_user) \
.add_step(send_welcome)
@staticmethod
def create_shopping_flow() -> BusinessFlow:
"""购物流程"""
flow = BusinessFlow("购物流程")
# 步骤1：用户登录
login = HttpRequestStep(
"用户登录",
"POST",
"/api/login",
json={
"username": "{username}",
"password": "{password}"
}
).add_validator(lambda r: r.assert_status_code(200)) \
.add_extractor("access_token", "$.data.access_token")
# 步骤2：浏览商品
browse_products = HttpRequestStep(
"浏览商品",
"GET",
"/api/products?category={category}&page=1",
headers={"Authorization": "Bearer {access_token}"}
).add_validator(lambda r: r.assert_status_code(200)) \
.add_extractor("product_id", "$.data.products[0].id")
# 步骤3：添加到购物车
add_to_cart = HttpRequestStep(
"添加到购物车",
"POST",
"/api/cart",
headers={"Authorization": "Bearer {access_token}"},
json={
"product_id": "{product_id}",
"quantity": 1
}
).add_validator(lambda r: r.assert_status_code(200))
# 步骤4：结算订单
checkout = HttpRequestStep(
"结算订单",
"POST",
"/api/orders",
headers={"Authorization": "Bearer {access_token}"},
json={
"payment_method": "credit_card",
"shipping_address": "{shipping_address}"
}
).add_validator(lambda r: r.assert_status_code(201)) \
.add_extractor("order_id", "$.data.order_id")
return flow.add_step(login) \
.add_step(browse_products) \
.add_step(add_to_cart) \
.add_step(checkout)3. 脚本引擎集成
# core/script_engine.py
"""
脚本引擎主类 - 整合所有组件
"""
from typing import Dict, Any, List
from .http_client import EnhancedHttpClient
from .business_flow import BusinessFlow
class ScriptEngine:
"""
压测脚本引擎
整合HTTP客户端、业务流程、数据管理等组件
"""
def __init__(self, user_instance):
self.user = user_instance
self.client = EnhancedHttpClient(user_instance.client)
self.flows: Dict[str, BusinessFlow] = {}
self.global_context = {}
# 设置默认请求头
self.client.set_default_headers({
'User-Agent': 'Locust Performance Test',
'Accept': 'application/json',
'Content-Type': 'application/json'
})
# 添加性能监控钩子
self.client.add_response_hook(self._performance_monitor_hook)
def register_flow(self, name: str, flow: BusinessFlow):
"""注册业务流程"""
self.flows[name] = flow
def execute_flow(self, flow_name: str, **context) -> Dict[str, Any]:
"""执行业务流程"""
if flow_name not in self.flows:
raise ValueError(f"未找到业务流程: {flow_name}")
flow = self.flows[flow_name]
# 合并全局上下文和局部上下文
execution_context = self.global_context.copy()
execution_context.update(context)
# 执行流程
result = flow.set_context(**execution_context).execute(self.client)
# 更新全局上下文
self.global_context.update(result)
return result
def _performance_monitor_hook(self, response):
"""性能监控钩子"""
if response.response_time > 5000: # 超过5秒
self.user.environment.events.user_error.fire(
user_instance=self.user,
exception=f"慢请求警告: {response.response_time}ms",
tb=""
)
return response
def set_global_context(self, **kwargs):
"""设置全局上下文"""
self.global_context.update(kwargs)
def get_context_value(self, key: str, default=None):
"""获取上下文值"""
return self.global_context.get(key, default)
# 使用示例
class ECommerceUser(BaseUser):
"""电商用户类 - 使用脚本引擎"""
def on_start(self):
super().on_start()
# 初始化脚本引擎
self.engine = ScriptEngine(self)
# 注册业务流程
self.engine.register_flow(
"user_registration",
ECommerceFlow.create_user_registration_flow()
)
self.engine.register_flow(
"shopping",
ECommerceFlow.create_shopping_flow()
)
# 设置全局上下文
self.engine.set_global_context(
username=f"test_user_{self.user_id}",
password="test123456",
email=f"test_{self.user_id}@example.com",
category="electronics",
shipping_address="测试地址123号"
)
@task(1)
def user_registration_flow(self):
"""用户注册流程测试"""
try:
result = self.engine.execute_flow("user_registration")
self.logger.info(f"用户注册成功: {result.get('user_id')}")
except Exception as e:
self.logger.error(f"用户注册失败: {e}")
@task(3)
def shopping_flow(self):
"""购物流程测试"""
try:
result = self.engine.execute_flow("shopping")
self.logger.info(f"购物完成，订单ID: {result.get('order_id')}")
except Exception as e:
self.logger.error(f"购物流程失败: {e}")高级特性：让脚本更智能
1. 智能重试机制
# core/retry_handler.py
"""
智能重试机制 - 让脚本更健壮
"""
import time
import random
from typing import Callable, Type, Tuple
class RetryHandler:
"""智能重试处理器"""
def __init__(self,
max_retries: int = 3,
backoff_factor: float = 1.0,
jitter: bool = True):
self.max_retries = max_retries
self.backoff_factor = backoff_factor
self.jitter = jitter
def retry(self,
func: Callable,
exceptions: Tuple[Type[Exception], ...] = (Exception,),
**kwargs):
"""执行重试逻辑"""
for attempt in range(self.max_retries + 1):
try:
return func(**kwargs)
except exceptions as e:
if attempt == self.max_retries:
raise
# 计算等待时间
wait_time = self.backoff_factor * (2 ** attempt)
if self.jitter:
wait_time += random.uniform(0, 1)
time.sleep(wait_time)
continue2. 响应缓存机制
# core/response_cache.py
"""
响应缓存机制 - 提升测试效率
"""
import hashlib
import json
from typing import Dict, Any, Optional
class ResponseCache:
"""响应缓存管理器"""
def __init__(self, max_size: int = 1000, ttl: int = 300):
self.max_size = max_size
self.ttl = ttl
self.cache: Dict[str, Dict[str, Any]] = {}
def _generate_key(self, method: str, url: str, **kwargs) -> str:
"""生成缓存键"""
key_data = {
'method': method,
'url': url,
'params': kwargs.get('params'),
'json': kwargs.get('json')
}
key_str = json.dumps(key_data, sort_keys=True)
return hashlib.md5(key_str.encode()).hexdigest()
def get(self, method: str, url: str, **kwargs) -> Optional[Any]:
"""获取缓存响应"""
key = self._generate_key(method, url, **kwargs)
if key in self.cache:
cache_entry = self.cache[key]
if time.time() - cache_entry['timestamp'] < self.ttl:
return cache_entry['response']
else:
del self.cache[key]
return None
def set(self, method: str, url: str, response: Any, **kwargs):
"""设置缓存响应"""
if len(self.cache) >= self.max_size:
# 清理最旧的缓存
oldest_key = min(self.cache.keys(),
key=lambda k: self.cache[k]['timestamp'])
del self.cache[oldest_key]
key = self._generate_key(method, url, **kwargs)
self.cache[key] = {
'response': response,
'timestamp': time.time()
}总结
压测脚本引擎的设计就像打造一把精密的工具，既要功能强大，又要使用简单。通过这篇文章，我们深入了解了：
- 设计理念：简单易用与功能强大的平衡
- HTTP客户端：增强的网络请求处理能力
- 业务流程建模：复杂场景的抽象和复用
- 脚本引擎集成：各组件的协调工作
- 高级特性：智能重试、缓存等优化功能
这个脚本引擎不仅仅是对Locust的简单封装，而是一个经过实战检验的、功能完整的业务场景测试解决方案。它让我们可以专注于业务逻辑的实现，而不用担心底层的技术细节。
下一篇文章，我们将探讨数据驱动测试与参数化设计，看看如何让测试数据管理变得更加灵活和智能。