easy-rag/tests/utils.py

241 lines
8.2 KiB
Python

"""
测试工具和辅助函数
"""
import asyncio
import time
import json
from datetime import datetime
from typing import Dict, List, Any
from pathlib import Path
class TestReporter:
"""测试报告生成器"""
def __init__(self, output_dir: str = "test_reports"):
self.output_dir = Path(output_dir)
self.output_dir.mkdir(exist_ok=True)
self.start_time = datetime.now()
def generate_report(self, results: Dict[str, Any], report_name: str = None):
"""生成测试报告"""
if not report_name:
report_name = f"test_report_{self.start_time.strftime('%Y%m%d_%H%M%S')}"
# 生成 Markdown 报告
md_content = self._generate_markdown_report(results)
md_file = self.output_dir / f"{report_name}.md"
with open(md_file, 'w', encoding='utf-8') as f:
f.write(md_content)
# 生成 JSON 报告
json_content = self._generate_json_report(results)
json_file = self.output_dir / f"{report_name}.json"
with open(json_file, 'w', encoding='utf-8') as f:
json.dump(json_content, f, indent=2, ensure_ascii=False)
return {
"markdown": str(md_file),
"json": str(json_file)
}
def _generate_markdown_report(self, results: Dict[str, Any]) -> str:
"""生成 Markdown 格式报告"""
report = f"""# RAG 系统测试报告
## 测试概览
- **测试时间**: {self.start_time.strftime('%Y-%m-%d %H:%M:%S')}
- **报告生成时间**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
## 测试结果汇总
"""
# 添加各项测试结果
for test_type, test_results in results.items():
if isinstance(test_results, list):
successful = sum(1 for r in test_results if isinstance(r, dict) and r.get('status_code') == 200)
total = len(test_results)
success_rate = (successful / total * 100) if total > 0 else 0
report += f"### {test_type.replace('_', ' ').title()}\n"
report += f"- 总请求数: {total}\n"
report += f"- 成功数: {successful}\n"
report += f"- 成功率: {success_rate:.1f}%\n"
if test_results:
avg_time = sum(r.get('response_time', 0) for r in test_results if isinstance(r, dict)) / len(test_results)
report += f"- 平均响应时间: {avg_time:.3f}\n"
report += "\n"
return report
def _generate_json_report(self, results: Dict[str, Any]) -> Dict[str, Any]:
"""生成 JSON 格式报告"""
return {
"test_info": {
"start_time": self.start_time.isoformat(),
"end_time": datetime.now().isoformat(),
"duration": (datetime.now() - self.start_time).total_seconds()
},
"results": results,
"summary": self._calculate_summary(results)
}
def _calculate_summary(self, results: Dict[str, Any]) -> Dict[str, Any]:
"""计算测试摘要"""
summary = {
"total_requests": 0,
"total_successful": 0,
"overall_success_rate": 0,
"test_types": len(results)
}
for test_results in results.values():
if isinstance(test_results, list):
summary["total_requests"] += len(test_results)
summary["total_successful"] += sum(
1 for r in test_results
if isinstance(r, dict) and r.get('status_code') == 200
)
if summary["total_requests"] > 0:
summary["overall_success_rate"] = (
summary["total_successful"] / summary["total_requests"] * 100
)
return summary
class TestDataGenerator:
"""测试数据生成器"""
@staticmethod
def generate_test_documents(count: int, base_content: str = None) -> List[Dict[str, str]]:
"""生成测试文档"""
if not base_content:
base_content = "这是一个测试文档,包含关于人工智能和机器学习的内容。"
documents = []
for i in range(count):
content = f"{base_content} 文档编号: {i+1}" + f"额外内容: {'AI技术' if i % 2 == 0 else 'ML算法'}" * 10
documents.append({
"content": content,
"filename": f"test_doc_{i+1:03d}.txt"
})
return documents
@staticmethod
def generate_test_questions(count: int) -> List[str]:
"""生成测试问题"""
base_questions = [
"什么是人工智能?",
"机器学习的应用有哪些?",
"深度学习和传统机器学习的区别?",
"自然语言处理的主要挑战?",
"计算机视觉技术的发展趋势?",
]
questions = []
for i in range(count):
base_q = base_questions[i % len(base_questions)]
questions.append(f"{base_q} (查询 {i+1})")
return questions
class PerformanceAnalyzer:
"""性能分析器"""
@staticmethod
def analyze_response_times(results: List[Dict[str, Any]]) -> Dict[str, float]:
"""分析响应时间"""
times = [r.get('response_time', 0) for r in results if isinstance(r, dict)]
if not times:
return {}
times.sort()
n = len(times)
return {
"min": min(times),
"max": max(times),
"avg": sum(times) / n,
"median": times[n // 2],
"p95": times[int(n * 0.95)] if n > 1 else times[0],
"p99": times[int(n * 0.99)] if n > 1 else times[0]
}
@staticmethod
def analyze_success_rates(results: List[Dict[str, Any]]) -> Dict[str, Any]:
"""分析成功率"""
total = len(results)
successful = sum(1 for r in results if isinstance(r, dict) and r.get('status_code') == 200)
return {
"total": total,
"successful": successful,
"failed": total - successful,
"success_rate": (successful / total * 100) if total > 0 else 0,
"failure_rate": ((total - successful) / total * 100) if total > 0 else 0
}
def format_duration(seconds: float) -> str:
"""格式化持续时间"""
if seconds < 1:
return f"{seconds * 1000:.1f}ms"
elif seconds < 60:
return f"{seconds:.2f}s"
else:
minutes = int(seconds // 60)
seconds = seconds % 60
return f"{minutes}m {seconds:.1f}s"
def print_test_summary(test_name: str, results: List[Dict[str, Any]]):
"""打印测试摘要"""
if not results:
print(f"{test_name}: 没有结果")
return
analyzer = PerformanceAnalyzer()
success_info = analyzer.analyze_success_rates(results)
time_info = analyzer.analyze_response_times(results)
print(f"{test_name}:")
print(f" - 成功率: {success_info['success_rate']:.1f}% ({success_info['successful']}/{success_info['total']})")
if time_info:
print(f" - 响应时间: 平均 {format_duration(time_info['avg'])}, "
f"最大 {format_duration(time_info['max'])}, "
f"P95 {format_duration(time_info['p95'])}")
async def wait_for_server(base_url: str, timeout: int = 30) -> bool:
"""等待服务器启动"""
import aiohttp
print(f"🔍 等待服务器启动 ({base_url})...")
async with aiohttp.ClientSession() as session:
for i in range(timeout):
try:
async with session.get(f"{base_url}/health", timeout=1) as response:
if response.status == 200:
print(f"✅ 服务器已启动 (耗时: {i+1}秒)")
return True
except:
pass
await asyncio.sleep(1)
if i % 5 == 4: # 每5秒显示一次等待状态
print(f"⏳ 仍在等待服务器启动... ({i+1}/{timeout})")
print(f"❌ 服务器启动超时 ({timeout}秒)")
return False