easy-rag/tests/quick_test.py

138 lines
4.3 KiB
Python
Raw Permalink Blame History

This file contains invisible Unicode characters

This file contains invisible Unicode characters that are indistinguishable to humans but may be processed differently by a computer. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python3
"""
快速测试运行脚本
用于快速验证系统功能
"""
import asyncio
import sys
import time
import os
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
from tests.test_concurrent import (
ConcurrentRAGTester,
test_concurrent_health_check,
test_concurrent_upload,
test_concurrent_chat,
)
async def quick_test():
"""快速测试所有主要功能"""
print("🚀 快速功能验证测试")
print("=" * 40)
try:
async with ConcurrentRAGTester() as tester:
# 1. 健康检查
print("1⃣ 健康检查...")
health = await tester.health_check()
if health["status_code"] != 200:
print(f"❌ 服务器不可用: {health}")
return False
print(f"✅ 服务器正常 (响应时间: {health['response_time']:.3f}s)")
# 2. 单个文档上传
print("\n2⃣ 文档上传测试...")
upload_result = await tester.upload_document(
"这是一个快速测试文档。包含关于人工智能和机器学习的基础知识。",
"quick_test.txt",
)
if upload_result["status_code"] != 200:
print(f"❌ 上传失败: {upload_result}")
return False
print(f"✅ 上传成功 (文档ID: {upload_result.get('document_id', 'N/A')})")
# 3. 等待处理
await asyncio.sleep(1)
# 4. 聊天测试
print("\n3⃣ 聊天功能测试...")
chat_result = await tester.chat_query("什么是人工智能?")
if chat_result["status_code"] != 200:
print(f"❌ 聊天失败: {chat_result}")
return False
print(
f"✅ 聊天成功 (处理时间: {chat_result.get('processing_time', 0):.2f}s)"
)
print(f" 回答长度: {len(chat_result.get('answer', ''))} 字符")
print(f" 来源数量: {chat_result.get('sources_count', 0)}")
# 5. 文档列表
print("\n4⃣ 文档列表测试...")
docs_result = await tester.get_documents()
if docs_result["status_code"] != 200:
print(f"❌ 获取文档列表失败: {docs_result}")
return False
doc_count = len(docs_result["data"])
print(f"✅ 文档列表获取成功 (文档数量: {doc_count})")
print("\n" + "=" * 40)
print("🎉 所有基础功能测试通过!")
return True
except Exception as e:
print(f"❌ 测试过程中发生错误: {e}")
return False
async def mini_concurrent_test():
"""迷你并发测试"""
print("\n🔥 迷你并发测试")
print("=" * 40)
try:
# 小规模并发测试
await test_concurrent_health_check(3)
await test_concurrent_upload(2)
await asyncio.sleep(1)
await test_concurrent_chat(3)
print("🎯 迷你并发测试完成!")
return True
except Exception as e:
print(f"❌ 并发测试失败: {e}")
return False
def main():
"""主函数"""
if len(sys.argv) > 1:
test_type = sys.argv[1].lower()
if test_type == "quick":
success = asyncio.run(quick_test())
elif test_type == "concurrent":
success = asyncio.run(mini_concurrent_test())
elif test_type == "both":
success1 = asyncio.run(quick_test())
success2 = asyncio.run(mini_concurrent_test()) if success1 else False
success = success1 and success2
else:
print("❌ 未知的测试类型")
print("用法: python quick_test.py [quick|concurrent|both]")
return
else:
# 默认运行所有测试
success1 = asyncio.run(quick_test())
success2 = asyncio.run(mini_concurrent_test()) if success1 else False
success = success1 and success2
if success:
print("\n✅ 所有测试通过!")
sys.exit(0)
else:
print("\n❌ 测试失败!")
sys.exit(1)
if __name__ == "__main__":
try:
main()
except KeyboardInterrupt:
print("\n⏹️ 测试被中断")
except Exception as e:
print(f"❌ 运行失败: {e}")
sys.exit(1)