103 lines
3.1 KiB
Python
103 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
测试MongoDB和BOS存储集成
|
||
"""
|
||
|
||
import os
|
||
import logging
|
||
from mongodb_manager import get_mongodb_manager
|
||
from baidu_bos_manager import get_bos_manager
|
||
|
||
logging.basicConfig(level=logging.INFO)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
def test_mongodb_connection():
|
||
"""测试MongoDB连接"""
|
||
try:
|
||
mongodb_mgr = get_mongodb_manager()
|
||
stats = mongodb_mgr.get_stats()
|
||
logger.info(f"✅ MongoDB连接成功，统计信息: {stats}")
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"❌ MongoDB连接失败: {e}")
|
||
return False
|
||
|
||
def test_bos_connection():
|
||
"""测试BOS连接"""
|
||
try:
|
||
bos_mgr = get_bos_manager()
|
||
objects = bos_mgr.list_objects(max_keys=5)
|
||
logger.info(f"✅ BOS连接成功，找到 {len(objects)} 个对象")
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"❌ BOS连接失败: {e}")
|
||
return False
|
||
|
||
def test_file_upload_workflow():
|
||
"""测试文件上传工作流"""
|
||
try:
|
||
# 创建测试文件
|
||
test_file = "/tmp/test_storage.txt"
|
||
with open(test_file, 'w', encoding='utf-8') as f:
|
||
f.write("这是一个测试文件，用于验证存储集成功能。")
|
||
|
||
# 获取管理器
|
||
mongodb_mgr = get_mongodb_manager()
|
||
bos_mgr = get_bos_manager()
|
||
|
||
# 上传到BOS
|
||
bos_result = bos_mgr.upload_file(test_file)
|
||
logger.info(f"✅ BOS上传成功: {bos_result['bos_key']}")
|
||
|
||
# 存储元数据到MongoDB
|
||
file_id = mongodb_mgr.store_file_metadata(
|
||
file_path=test_file,
|
||
file_type="text",
|
||
bos_key=bos_result["bos_key"],
|
||
additional_info={
|
||
"test": True,
|
||
"bos_url": bos_result["url"]
|
||
}
|
||
)
|
||
logger.info(f"✅ MongoDB元数据存储成功: {file_id}")
|
||
|
||
# 存储向量元数据
|
||
mongodb_mgr.store_vector_metadata(
|
||
file_id=file_id,
|
||
vector_type="text_vector",
|
||
vdb_id=bos_result["bos_key"],
|
||
vector_info={"test": True}
|
||
)
|
||
logger.info("✅ 向量元数据存储成功")
|
||
|
||
# 清理测试文件
|
||
os.remove(test_file)
|
||
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 文件上传工作流测试失败: {e}")
|
||
return False
|
||
|
||
if __name__ == "__main__":
|
||
logger.info("🚀 开始存储集成测试...")
|
||
|
||
# 测试MongoDB连接
|
||
mongodb_ok = test_mongodb_connection()
|
||
|
||
# 测试BOS连接
|
||
bos_ok = test_bos_connection()
|
||
|
||
# 测试完整工作流
|
||
workflow_ok = test_file_upload_workflow()
|
||
|
||
# 总结
|
||
if mongodb_ok and bos_ok and workflow_ok:
|
||
logger.info("🎉 所有存储集成测试通过！")
|
||
else:
|
||
logger.error("❌ 存储集成测试失败")
|
||
logger.error(f"MongoDB: {'✅' if mongodb_ok else '❌'}")
|
||
logger.error(f"BOS: {'✅' if bos_ok else '❌'}")
|
||
logger.error(f"工作流: {'✅' if workflow_ok else '❌'}")
|