276 lines
9.5 KiB
Python
276 lines
9.5 KiB
Python
"""
|
||
测试知识库路由的变更
|
||
文件: server/routers/knowledge_router.py
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
project_root = Path(__file__).parent.parent
|
||
sys.path.insert(0, str(project_root))
|
||
|
||
import pytest
|
||
from unittest.mock import Mock, AsyncMock, patch, MagicMock
|
||
from fastapi import UploadFile
|
||
from io import BytesIO
|
||
|
||
|
||
class TestKnowledgeRouterChanges:
|
||
"""测试 knowledge_router.py 的变更"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_add_documents_minio_upload(self):
|
||
"""测试上传文档后自动上传到 MinIO"""
|
||
# 模拟成功处理的文件项
|
||
success_items = [
|
||
{
|
||
"status": "done",
|
||
"path": "/tmp/test_file.pdf",
|
||
"filename": "test_file.pdf",
|
||
"file_type": "application/pdf",
|
||
}
|
||
]
|
||
|
||
# 模拟 aupload_file_to_minio 函数
|
||
with patch("server.routers.knowledge_router.aupload_file_to_minio") as mock_upload:
|
||
mock_upload.return_value = "https://minio.example.com/ref-test-db/test_file.pdf"
|
||
|
||
# 模拟 aiofiles.open
|
||
with patch("server.routers.knowledge_router.aiofiles.open") as mock_open:
|
||
mock_file = AsyncMock()
|
||
mock_file.read = AsyncMock(return_value=b"PDF content")
|
||
mock_open.return_value.__aenter__.return_value = mock_file
|
||
|
||
# 模拟上传逻辑
|
||
for success_item in success_items:
|
||
# 读取文件
|
||
async with mock_open(success_item["path"], "rb") as f:
|
||
file_bytes = await f.read()
|
||
|
||
# 上传到 MinIO
|
||
refdb = "test_db".replace("_", "-")
|
||
url = await mock_upload(
|
||
f"ref-{refdb}",
|
||
success_item["filename"],
|
||
file_bytes,
|
||
success_item["file_type"],
|
||
)
|
||
|
||
assert url is not None
|
||
assert "ref-test-db" in url
|
||
print(f"✓ MinIO 上传模拟成功: {url}")
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_delete_document_minio_cleanup(self):
|
||
"""测试删除文档时同时删除 MinIO 中的文件"""
|
||
db_id = "test_database"
|
||
doc_id = "doc_12345"
|
||
file_name = "test_document.pdf"
|
||
|
||
# 模拟 knowledge_base.get_file_basic_info
|
||
with patch("server.routers.knowledge_router.knowledge_base") as mock_kb:
|
||
mock_kb.get_file_basic_info = AsyncMock(
|
||
return_value={"meta": {"filename": file_name}}
|
||
)
|
||
mock_kb.delete_file = AsyncMock()
|
||
|
||
# 模拟 MinIO 客户端
|
||
with patch("server.routers.knowledge_router.get_minio_client") as mock_get_client:
|
||
mock_client = Mock()
|
||
mock_client.adelete_file = AsyncMock(return_value=True)
|
||
mock_get_client.return_value = mock_client
|
||
|
||
# 执行删除逻辑
|
||
file_meta_info = await mock_kb.get_file_basic_info(db_id, doc_id)
|
||
file_name_result = file_meta_info.get("meta", {}).get("filename")
|
||
|
||
await mock_client.adelete_file("ref-" + db_id.replace("_", "-"), file_name_result)
|
||
await mock_kb.delete_file(db_id, doc_id)
|
||
|
||
# 验证调用
|
||
mock_kb.get_file_basic_info.assert_called_once_with(db_id, doc_id)
|
||
mock_client.adelete_file.assert_called_once_with(
|
||
"ref-test-database", file_name
|
||
)
|
||
mock_kb.delete_file.assert_called_once_with(db_id, doc_id)
|
||
print(f"✓ 删除文档时 MinIO 清理逻辑验证成功")
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_download_document_streaming_response(self):
|
||
"""测试文档下载的流式响应"""
|
||
import asyncio
|
||
|
||
# 模拟 MinIO 响应对象
|
||
class MockMinioResponse:
|
||
def __init__(self, data):
|
||
self.data = data
|
||
self.position = 0
|
||
|
||
def read(self, size):
|
||
"""同步读取方法"""
|
||
if self.position >= len(self.data):
|
||
return b""
|
||
chunk = self.data[self.position : self.position + size]
|
||
self.position += size
|
||
return chunk
|
||
|
||
def close(self):
|
||
pass
|
||
|
||
def release_conn(self):
|
||
pass
|
||
|
||
test_data = b"A" * 20000 # 20KB 测试数据
|
||
mock_response = MockMinioResponse(test_data)
|
||
|
||
# 模拟流式生成器
|
||
async def minio_stream():
|
||
try:
|
||
while True:
|
||
chunk = await asyncio.to_thread(mock_response.read, 8192)
|
||
if not chunk:
|
||
break
|
||
yield chunk
|
||
finally:
|
||
mock_response.close()
|
||
mock_response.release_conn()
|
||
|
||
# 收集所有块
|
||
chunks = []
|
||
async for chunk in minio_stream():
|
||
chunks.append(chunk)
|
||
|
||
# 验证
|
||
result_data = b"".join(chunks)
|
||
assert len(result_data) == len(test_data)
|
||
assert result_data == test_data
|
||
print(f"✓ 流式下载响应验证成功（数据大小: {len(result_data)} 字节，{len(chunks)} 个块）")
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_upload_file_hash_calculation(self):
|
||
"""测试文件上传时的异步哈希计算"""
|
||
from src.knowledge.utils.kb_utils import calculate_content_hash
|
||
|
||
# 模拟上传的文件内容
|
||
file_content = b"Test file content for hashing"
|
||
|
||
# 异步计算哈希
|
||
content_hash = await calculate_content_hash(file_content)
|
||
|
||
assert content_hash is not None
|
||
assert len(content_hash) == 64 # SHA-256
|
||
print(f"✓ 上传文件哈希计算成功: {content_hash}")
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_file_existence_check(self):
|
||
"""测试文件存在性检查的异步调用"""
|
||
db_id = "test_db"
|
||
content_hash = "a" * 64
|
||
|
||
# 模拟 knowledge_base
|
||
with patch("server.routers.knowledge_router.knowledge_base") as mock_kb:
|
||
mock_kb.file_existed_in_db = AsyncMock(return_value=False)
|
||
|
||
# 执行检查
|
||
file_exists = await mock_kb.file_existed_in_db(db_id, content_hash)
|
||
|
||
assert file_exists is False
|
||
mock_kb.file_existed_in_db.assert_called_once_with(db_id, content_hash)
|
||
print(f"✓ 异步文件存在性检查验证成功")
|
||
|
||
def test_upload_file_fixed_salt_logic(self):
|
||
"""测试上传文件时的固定 salt 逻辑"""
|
||
from src.utils import hashstr
|
||
|
||
# 测试固定 salt 的文件命名
|
||
basename = "test_document"
|
||
ext = ".pdf"
|
||
|
||
# 使用固定 salt
|
||
filename1 = f"{basename}_{hashstr(basename, 4, with_salt=True, salt='fixed_salt')}{ext}".lower()
|
||
filename2 = f"{basename}_{hashstr(basename, 4, with_salt=True, salt='fixed_salt')}{ext}".lower()
|
||
|
||
# 验证相同名称生成相同的哈希（因为 salt 固定）
|
||
assert filename1 == filename2
|
||
print(f"✓ 固定 salt 文件命名验证成功: {filename1}")
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_streaming_response_error_handling(self):
|
||
"""测试流式响应的错误处理"""
|
||
import asyncio
|
||
|
||
class MockErrorResponse:
|
||
def read(self, size):
|
||
raise Exception("MinIO read error")
|
||
|
||
def close(self):
|
||
pass
|
||
|
||
def release_conn(self):
|
||
pass
|
||
|
||
mock_response = MockErrorResponse()
|
||
|
||
async def minio_stream():
|
||
try:
|
||
while True:
|
||
chunk = await asyncio.to_thread(mock_response.read, 8192)
|
||
if not chunk:
|
||
break
|
||
yield chunk
|
||
except Exception as e:
|
||
print(f"✓ 错误被正确捕获: {e}")
|
||
raise
|
||
finally:
|
||
mock_response.close()
|
||
mock_response.release_conn()
|
||
|
||
# 验证错误处理
|
||
with pytest.raises(Exception, match="MinIO read error"):
|
||
async for chunk in minio_stream():
|
||
pass
|
||
|
||
print(f"✓ 流式响应错误处理验证成功")
|
||
|
||
|
||
class TestMinIOIntegration:
|
||
"""测试 MinIO 集成逻辑"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_bucket_naming_convention(self):
|
||
"""测试 bucket 命名约定（ref-{db_id}）"""
|
||
db_id = "test_knowledge_base"
|
||
expected_bucket = "ref-test-knowledge-base" # 下划线替换为连字符
|
||
|
||
refdb = db_id.replace("_", "-")
|
||
actual_bucket = f"ref-{refdb}"
|
||
|
||
assert actual_bucket == expected_bucket
|
||
print(f"✓ Bucket 命名约定正确: {db_id} -> {actual_bucket}")
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_concurrent_minio_uploads(self):
|
||
"""测试并发上传到 MinIO"""
|
||
import asyncio
|
||
|
||
# 模拟多个文件的并发上传
|
||
async def mock_upload(file_id):
|
||
await asyncio.sleep(0.1) # 模拟上传延迟
|
||
return f"https://minio.example.com/bucket/file_{file_id}"
|
||
|
||
# 并发上传 5 个文件
|
||
tasks = [mock_upload(i) for i in range(5)]
|
||
results = await asyncio.gather(*tasks)
|
||
|
||
assert len(results) == 5
|
||
for i, url in enumerate(results):
|
||
assert f"file_{i}" in url
|
||
|
||
print(f"✓ 并发 MinIO 上传模拟成功，处理了 {len(results)} 个文件")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
exit_code = pytest.main([__file__, "-v", "-s", "--tb=short", "--color=yes"])
|
||
sys.exit(exit_code)
|