feat: 新增多格式文件上传入库与认证体系
- 新增 JWT 认证模块,支持登录/注册/用户管理 - 新增文件上传接口,支持 .txt/.md/.html/.pdf/.docx 等格式解析入库 - 新增检索结果 AI 总结功能 - 新增文本去重缓存机制 - 新增全局认证夹具简化测试 - 新增配置项与环境变量支持 - 完善文档与测试覆盖
This commit is contained in:
@@ -0,0 +1,27 @@
|
||||
"""pytest 全局夹具:覆盖 JWT 认证依赖,让现有 API 测试默认以 admin 身份运行
|
||||
|
||||
业务接口(search/document/knowledge)已加 Depends(get_current_user)、
|
||||
DELETE /documents 加了 Depends(require_admin)。这里通过 autouse 夹具把两个依赖
|
||||
统一替换为返回固定 admin AuthUser 的 lambda,使现有 API 测试无需改动即可通过认证。
|
||||
单个测试需要走真实认证逻辑时(如 tests/test_auth.py),可在测试函数内
|
||||
pop 掉对应 override,autouse fixture yield 后会统一 clear。
|
||||
"""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from app.core.auth import get_current_user, require_admin
|
||||
from app.main import app
|
||||
from app.models.auth import AuthUser
|
||||
|
||||
TEST_USER = AuthUser(username="testuser", role="admin", created_at=datetime.now(UTC))
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def override_auth():
|
||||
"""所有测试默认以 admin 身份运行;测试结束清理 dependency_overrides"""
|
||||
app.dependency_overrides[get_current_user] = lambda: TEST_USER
|
||||
app.dependency_overrides[require_admin] = lambda: TEST_USER
|
||||
yield
|
||||
app.dependency_overrides.clear()
|
||||
@@ -0,0 +1,293 @@
|
||||
"""认证 API 与核心函数的单元测试(mock UserStore,不依赖真实 Redis)
|
||||
|
||||
覆盖:
|
||||
- POST /auth/login:成功 / 密码错误 / 用户不存在
|
||||
- POST /auth/register:成功 / 用户已存在 / 注册关闭
|
||||
- GET /auth/me:有效 token / 无 token / 无效 token(需走真实 get_current_user)
|
||||
- require_admin:非 admin 抛 FORBIDDEN(直接测函数)
|
||||
- create_access_token + decode_token 往返一致
|
||||
- hash_password + verify_password 正确 / 错误
|
||||
|
||||
UserStore 的 authenticate/create 在 FakeUserStore 中 mock,避免依赖真实 Redis。
|
||||
"""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.api.response import ApiError
|
||||
from app.api.v1 import auth as auth_module
|
||||
from app.config import settings
|
||||
from app.core import auth as auth_core
|
||||
from app.core.auth import (
|
||||
ERR_BAD_CREDENTIALS,
|
||||
ERR_FORBIDDEN,
|
||||
ERR_REGISTER_DISABLED,
|
||||
ERR_TOKEN_INVALID,
|
||||
ERR_UNAUTHORIZED,
|
||||
ERR_USER_EXISTS,
|
||||
create_access_token,
|
||||
decode_token,
|
||||
get_current_user,
|
||||
hash_password,
|
||||
require_admin,
|
||||
verify_password,
|
||||
)
|
||||
from app.main import app
|
||||
from app.models.auth import AuthUser, StoredUser
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
class FakeUserStore:
|
||||
"""假 UserStore:按预置数据返回结果或抛 ApiError,记录调用
|
||||
|
||||
authenticate/create 使用预置 hashed_password,不调用真实 hash_password。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
user: StoredUser | None = None,
|
||||
exists: bool = False,
|
||||
create_error: ApiError | None = None,
|
||||
auth_fail: bool = False,
|
||||
) -> None:
|
||||
self.user = user
|
||||
self._exists = exists
|
||||
self.create_error = create_error
|
||||
self.auth_fail = auth_fail
|
||||
self.create_calls: list[tuple[str, str, str]] = []
|
||||
|
||||
async def get(self, username: str) -> StoredUser | None:
|
||||
if self.user is not None and self.user.username == username:
|
||||
return self.user
|
||||
return None
|
||||
|
||||
async def exists(self, username: str) -> bool:
|
||||
return self._exists
|
||||
|
||||
async def create(self, username: str, password: str, role: str = "user") -> StoredUser:
|
||||
if self.create_error is not None:
|
||||
raise self.create_error
|
||||
user = StoredUser(
|
||||
username=username,
|
||||
role=role,
|
||||
created_at=_now(),
|
||||
hashed_password="fake-hash",
|
||||
)
|
||||
self.create_calls.append((username, password, role))
|
||||
return user
|
||||
|
||||
async def authenticate(self, username: str, password: str) -> StoredUser:
|
||||
if self.auth_fail or self.user is None or self.user.username != username:
|
||||
raise ApiError(ERR_BAD_CREDENTIALS, "用户名或密码错误")
|
||||
return self.user
|
||||
|
||||
|
||||
class NullUserStore:
|
||||
"""恒返回 None 的空存储:用于 /auth/me 测试,让 get_current_user 降级用 JWT payload"""
|
||||
|
||||
async def get(self, username: str) -> StoredUser | None:
|
||||
return None
|
||||
|
||||
async def exists(self, username: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def _install_store(monkeypatch: pytest.MonkeyPatch, store: FakeUserStore) -> None:
|
||||
"""将 auth 路由模块的 get_user_store 替换为假存储"""
|
||||
monkeypatch.setattr(auth_module, "get_user_store", lambda: store)
|
||||
|
||||
|
||||
def _make_stored_user(username: str = "alice", role: str = "user") -> StoredUser:
|
||||
return StoredUser(username=username, role=role, created_at=_now(), hashed_password="fake-hash")
|
||||
|
||||
|
||||
class TestAuthLoginApi:
|
||||
"""POST /auth/login"""
|
||||
|
||||
def test_login_success(self, monkeypatch: pytest.MonkeyPatch):
|
||||
user = _make_stored_user(username="alice", role="admin")
|
||||
_install_store(monkeypatch, FakeUserStore(user=user))
|
||||
|
||||
client = TestClient(app)
|
||||
resp = client.post("/api/v1/auth/login", json={"username": "alice", "password": "secret123"})
|
||||
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert data["access_token"]
|
||||
assert data["token_type"] == "bearer"
|
||||
assert data["expires_in"] == settings.jwt_expire_minutes * 60
|
||||
assert data["user"]["username"] == "alice"
|
||||
assert data["user"]["role"] == "admin"
|
||||
assert "created_at" in data["user"]
|
||||
|
||||
def test_login_wrong_password(self, monkeypatch: pytest.MonkeyPatch):
|
||||
user = _make_stored_user(username="alice")
|
||||
# authenticate 失败(密码不匹配)
|
||||
_install_store(monkeypatch, FakeUserStore(user=user, auth_fail=True))
|
||||
|
||||
client = TestClient(app)
|
||||
resp = client.post("/api/v1/auth/login", json={"username": "alice", "password": "WRONG"})
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == ERR_BAD_CREDENTIALS
|
||||
assert body["data"] is None
|
||||
|
||||
def test_login_user_not_found(self, monkeypatch: pytest.MonkeyPatch):
|
||||
_install_store(monkeypatch, FakeUserStore(user=None))
|
||||
|
||||
client = TestClient(app)
|
||||
resp = client.post("/api/v1/auth/login", json={"username": "nobody", "password": "whatever"})
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == ERR_BAD_CREDENTIALS
|
||||
|
||||
|
||||
class TestAuthRegisterApi:
|
||||
"""POST /auth/register"""
|
||||
|
||||
def test_register_success(self, monkeypatch: pytest.MonkeyPatch):
|
||||
store = FakeUserStore(user=None, exists=False)
|
||||
_install_store(monkeypatch, store)
|
||||
|
||||
client = TestClient(app)
|
||||
resp = client.post(
|
||||
"/api/v1/auth/register", json={"username": "newbie", "password": "passwd123"}
|
||||
)
|
||||
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert data["access_token"]
|
||||
assert data["user"]["username"] == "newbie"
|
||||
assert data["user"]["role"] == "user"
|
||||
assert len(store.create_calls) == 1
|
||||
assert store.create_calls[0][0] == "newbie"
|
||||
assert store.create_calls[0][2] == "user"
|
||||
|
||||
def test_register_user_exists(self, monkeypatch: pytest.MonkeyPatch):
|
||||
store = FakeUserStore(
|
||||
exists=True,
|
||||
create_error=ApiError(ERR_USER_EXISTS, "用户名已存在: newbie"),
|
||||
)
|
||||
_install_store(monkeypatch, store)
|
||||
|
||||
client = TestClient(app)
|
||||
resp = client.post(
|
||||
"/api/v1/auth/register", json={"username": "newbie", "password": "passwd123"}
|
||||
)
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == ERR_USER_EXISTS
|
||||
assert body["data"] is None
|
||||
|
||||
def test_register_disabled(self, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(settings, "auth_register_enabled", False)
|
||||
_install_store(monkeypatch, FakeUserStore())
|
||||
|
||||
client = TestClient(app)
|
||||
resp = client.post(
|
||||
"/api/v1/auth/register", json={"username": "newbie", "password": "passwd123"}
|
||||
)
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == ERR_REGISTER_DISABLED
|
||||
assert body["data"] is None
|
||||
|
||||
|
||||
class TestAuthMeApi:
|
||||
"""GET /auth/me:需走真实 get_current_user,测试内清除 override"""
|
||||
|
||||
def test_me_with_valid_token(self, monkeypatch: pytest.MonkeyPatch):
|
||||
# 清除 override,让 get_current_user 走真实认证
|
||||
app.dependency_overrides.pop(get_current_user, None)
|
||||
# mock get_user_store 返回空存储(模拟 Redis 无该用户,降级用 JWT payload)
|
||||
monkeypatch.setattr(auth_core, "get_user_store", lambda: NullUserStore())
|
||||
|
||||
token, _ = create_access_token("alice", "admin")
|
||||
client = TestClient(app)
|
||||
resp = client.get("/api/v1/auth/me", headers={"Authorization": f"Bearer {token}"})
|
||||
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
assert body["data"]["username"] == "alice"
|
||||
assert body["data"]["role"] == "admin"
|
||||
|
||||
def test_me_without_token(self, monkeypatch: pytest.MonkeyPatch):
|
||||
app.dependency_overrides.pop(get_current_user, None)
|
||||
monkeypatch.setattr(auth_core, "get_user_store", lambda: NullUserStore())
|
||||
|
||||
client = TestClient(app)
|
||||
resp = client.get("/api/v1/auth/me")
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == ERR_UNAUTHORIZED
|
||||
assert body["data"] is None
|
||||
|
||||
def test_me_with_invalid_token(self, monkeypatch: pytest.MonkeyPatch):
|
||||
app.dependency_overrides.pop(get_current_user, None)
|
||||
monkeypatch.setattr(auth_core, "get_user_store", lambda: NullUserStore())
|
||||
|
||||
client = TestClient(app)
|
||||
resp = client.get("/api/v1/auth/me", headers={"Authorization": "Bearer not-a-jwt"})
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == ERR_TOKEN_INVALID
|
||||
assert body["data"] is None
|
||||
|
||||
|
||||
class TestRequireAdmin:
|
||||
"""require_admin:admin 通过、非 admin 抛 FORBIDDEN(直接测函数,不经 API)"""
|
||||
|
||||
async def test_admin_passes(self):
|
||||
admin = AuthUser(username="alice", role="admin", created_at=_now())
|
||||
result = await require_admin(user=admin)
|
||||
assert result.role == "admin"
|
||||
|
||||
async def test_non_admin_raises_forbidden(self):
|
||||
normal = AuthUser(username="bob", role="user", created_at=_now())
|
||||
with pytest.raises(ApiError) as exc:
|
||||
await require_admin(user=normal)
|
||||
assert exc.value.code == ERR_FORBIDDEN
|
||||
|
||||
|
||||
class TestAuthCoreFunctions:
|
||||
"""核心函数:token 往返 / 密码哈希"""
|
||||
|
||||
def test_create_and_decode_token_roundtrip(self):
|
||||
token, expires_in = create_access_token("alice", "admin")
|
||||
assert expires_in == settings.jwt_expire_minutes * 60
|
||||
payload = decode_token(token)
|
||||
assert payload["sub"] == "alice"
|
||||
assert payload["role"] == "admin"
|
||||
assert "iat" in payload and "exp" in payload
|
||||
|
||||
def test_decode_invalid_token_raises(self):
|
||||
with pytest.raises(ApiError) as exc:
|
||||
decode_token("not-a-jwt")
|
||||
assert exc.value.code == ERR_TOKEN_INVALID
|
||||
|
||||
def test_hash_password_not_plaintext(self):
|
||||
hashed = hash_password("my-secret")
|
||||
assert hashed != "my-secret"
|
||||
assert hashed.startswith("$2") # bcrypt 哈希前缀
|
||||
|
||||
def test_verify_password_correct(self):
|
||||
hashed = hash_password("my-secret")
|
||||
assert verify_password("my-secret", hashed) is True
|
||||
|
||||
def test_verify_password_wrong(self):
|
||||
hashed = hash_password("my-secret")
|
||||
assert verify_password("wrong-password", hashed) is False
|
||||
|
||||
def test_verify_password_garbage_hash_returns_false(self):
|
||||
# 非法 hash 会被 bcrypt 拒绝,verify_password 捕获异常返回 False
|
||||
assert verify_password("any", "not-a-valid-hash") is False
|
||||
@@ -0,0 +1,313 @@
|
||||
"""POST /api/v1/documents/upload 端点测试(TestClient + FakeManager,不真实联网)"""
|
||||
|
||||
import io
|
||||
import re
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.api.v1 import document as document_module
|
||||
from app.main import app
|
||||
from app.models.document import DocumentInput
|
||||
from app.services.qdrant import QdrantService
|
||||
|
||||
|
||||
class FakeManager:
|
||||
def __init__(self, task_id: str = "task-upload-1") -> None:
|
||||
self.task_id = task_id
|
||||
self.submitted: list[DocumentInput] = []
|
||||
|
||||
async def submit(self, doc: DocumentInput) -> str:
|
||||
self.submitted.append(doc)
|
||||
return self.task_id
|
||||
|
||||
async def get(self, task_id: str) -> dict[str, Any] | None:
|
||||
return None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Iterator[TestClient]:
|
||||
async def _noop_ensure_collections(self: QdrantService) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(QdrantService, "ensure_collections", _noop_ensure_collections)
|
||||
monkeypatch.setattr(document_module.settings, "upload_dir", str(tmp_path / "uploads"))
|
||||
with TestClient(app) as test_client:
|
||||
yield test_client
|
||||
|
||||
|
||||
def _inject_manager(monkeypatch: pytest.MonkeyPatch, manager: FakeManager) -> None:
|
||||
monkeypatch.setattr(document_module, "_get_task_manager", lambda: manager)
|
||||
|
||||
|
||||
def _make_minimal_pdf(text: str = "Hello PDF World") -> bytes:
|
||||
content_stream = f"BT /F1 24 Tf 100 700 Td ({text}) Tj ET".encode("latin-1")
|
||||
content_obj = b"<< /Length " + str(len(content_stream)).encode() + b" >>\nstream\n" + content_stream + b"\nendstream"
|
||||
return (
|
||||
b"%PDF-1.0\n"
|
||||
b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n"
|
||||
b"2 0 obj\n<< /Type /Pages /Kids [3 0 R] /Count 1 >>\nendobj\n"
|
||||
b"3 0 obj\n<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] "
|
||||
b"/Resources << /Font << /F1 4 0 R >> >> /Contents 5 0 R >>\nendobj\n"
|
||||
b"4 0 obj\n<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>\nendobj\n"
|
||||
b"5 0 obj\n" + content_obj + b"\nendobj\n"
|
||||
b"xref\n0 6\n"
|
||||
b"0000000000 65535 f\n"
|
||||
b"0000000010 00000 n\n"
|
||||
b"0000000059 00000 n\n"
|
||||
b"0000000115 00000 n\n"
|
||||
b"0000000241 00000 n\n"
|
||||
b"0000000316 00000 n\n"
|
||||
b"trailer\n<< /Size 6 /Root 1 0 R >>\nstartxref\n414\n%%EOF\n"
|
||||
)
|
||||
|
||||
|
||||
def _make_docx(text_lines: list[str]) -> bytes:
|
||||
from docx import Document # type: ignore[import-untyped]
|
||||
|
||||
document = Document()
|
||||
for line in text_lines:
|
||||
document.add_paragraph(line)
|
||||
buf = io.BytesIO()
|
||||
document.save(buf)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def test_upload_md_returns_202_and_saves_file(
|
||||
client: TestClient, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""md 上传:202 + 落盘 + metadata 含原文件信息"""
|
||||
manager = FakeManager(task_id="abc-upload")
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
content = "# Hello\n\nThis is a markdown file.".encode("utf-8")
|
||||
resp = client.post(
|
||||
"/api/v1/documents/upload",
|
||||
files={"file": ("notes.md", content, "text/markdown")},
|
||||
data={"title": "我的笔记", "source": "manual"},
|
||||
)
|
||||
|
||||
assert resp.status_code == 202
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
data = body["data"]
|
||||
assert data["task_id"] == "abc-upload"
|
||||
assert data["status"] == "pending"
|
||||
|
||||
saved_path = data["saved_path"]
|
||||
assert saved_path
|
||||
assert Path(saved_path).read_bytes() == content
|
||||
|
||||
assert len(manager.submitted) == 1
|
||||
doc = manager.submitted[0]
|
||||
assert doc.text == content.decode("utf-8")
|
||||
assert doc.title == "我的笔记"
|
||||
assert doc.source == "manual"
|
||||
assert doc.metadata["original_filename"] == "notes.md"
|
||||
assert doc.metadata["original_size_bytes"] == str(len(content))
|
||||
assert doc.metadata["raw_file_path"] == saved_path
|
||||
|
||||
|
||||
def test_upload_txt_default_title_and_source(
|
||||
client: TestClient, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""txt 上传未传 title/source:默认取文件名 stem 与 file:{原文件名}"""
|
||||
manager = FakeManager()
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
resp = client.post(
|
||||
"/api/v1/documents/upload",
|
||||
files={"file": ("readme.txt", b"plain text body", "text/plain")},
|
||||
)
|
||||
|
||||
assert resp.status_code == 202
|
||||
assert len(manager.submitted) == 1
|
||||
doc = manager.submitted[0]
|
||||
assert doc.title == "readme"
|
||||
assert doc.source == "file:readme.txt"
|
||||
|
||||
|
||||
def test_upload_pdf_extracts_text_and_returns_202(
|
||||
client: TestClient, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""pdf 上传:提取文本并返回 202"""
|
||||
manager = FakeManager()
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
content = _make_minimal_pdf("Hello PDF World")
|
||||
resp = client.post(
|
||||
"/api/v1/documents/upload",
|
||||
files={"file": ("doc.pdf", content, "application/pdf")},
|
||||
)
|
||||
|
||||
assert resp.status_code == 202
|
||||
body = resp.json()
|
||||
assert body["code"] == 0
|
||||
assert body["data"]["saved_path"]
|
||||
assert len(manager.submitted) == 1
|
||||
assert "Hello PDF World" in manager.submitted[0].text
|
||||
|
||||
|
||||
def test_upload_docx_extracts_text(
|
||||
client: TestClient, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""docx 上传:提取段落文本"""
|
||||
manager = FakeManager()
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
content = _make_docx(["第一段落", "第二段落"])
|
||||
resp = client.post(
|
||||
"/api/v1/documents/upload",
|
||||
files={"file": ("doc.docx", content, "application/vnd.openxmlformats-officedocument.wordprocessingml.document")},
|
||||
)
|
||||
|
||||
assert resp.status_code == 202
|
||||
assert len(manager.submitted) == 1
|
||||
text = manager.submitted[0].text
|
||||
assert "第一段落" in text
|
||||
assert "第二段落" in text
|
||||
|
||||
|
||||
def test_upload_metadata_json_is_parsed(
|
||||
client: TestClient, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""metadata JSON 字符串被解析并入 metadata 字典"""
|
||||
manager = FakeManager()
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
resp = client.post(
|
||||
"/api/v1/documents/upload",
|
||||
files={"file": ("x.txt", b"content", "text/plain")},
|
||||
data={"metadata": '{"author": "alice", "team": "backend"}'},
|
||||
)
|
||||
|
||||
assert resp.status_code == 202
|
||||
assert len(manager.submitted) == 1
|
||||
doc = manager.submitted[0]
|
||||
assert doc.metadata["author"] == "alice"
|
||||
assert doc.metadata["team"] == "backend"
|
||||
assert "raw_file_path" in doc.metadata
|
||||
assert doc.metadata["original_filename"] == "x.txt"
|
||||
|
||||
|
||||
def test_upload_rejects_unsupported_extension(
|
||||
client: TestClient, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""不支持的扩展名:code=1001,message 含扩展名,未提交任务"""
|
||||
manager = FakeManager()
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
resp = client.post(
|
||||
"/api/v1/documents/upload",
|
||||
files={"file": ("data.xlsx", b"binary content", "application/octet-stream")},
|
||||
)
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 1001
|
||||
assert ".xlsx" in body["message"]
|
||||
assert manager.submitted == []
|
||||
|
||||
|
||||
def test_upload_rejects_oversized_file(
|
||||
client: TestClient, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""超过大小上限:code=1001,message 含'大小上限',未提交任务"""
|
||||
manager = FakeManager()
|
||||
_inject_manager(monkeypatch, manager)
|
||||
monkeypatch.setattr(document_module.settings, "upload_max_size_mb", 1)
|
||||
|
||||
big_content = b"x" * (2 * 1024 * 1024)
|
||||
resp = client.post(
|
||||
"/api/v1/documents/upload",
|
||||
files={"file": ("big.txt", big_content, "text/plain")},
|
||||
)
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 1001
|
||||
assert "大小上限" in body["message"]
|
||||
assert manager.submitted == []
|
||||
|
||||
|
||||
def test_upload_rejects_empty_text_after_parse(
|
||||
client: TestClient, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""解析后文本为空:code=1001,message 含'无法从文件提取文本',未提交任务"""
|
||||
manager = FakeManager()
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
resp = client.post(
|
||||
"/api/v1/documents/upload",
|
||||
files={"file": ("blank.txt", b" \n\t ", "text/plain")},
|
||||
)
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 1001
|
||||
assert "无法从文件提取文本" in body["message"]
|
||||
assert manager.submitted == []
|
||||
|
||||
|
||||
def test_upload_rejects_corrupted_pdf(
|
||||
client: TestClient, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""损坏的 PDF:code=1001,message 含'文件解析失败',未提交任务"""
|
||||
manager = FakeManager()
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
resp = client.post(
|
||||
"/api/v1/documents/upload",
|
||||
files={"file": ("bad.pdf", b"not a real pdf", "application/pdf")},
|
||||
)
|
||||
|
||||
body = resp.json()
|
||||
assert body["code"] == 1001
|
||||
assert "文件解析失败" in body["message"]
|
||||
assert manager.submitted == []
|
||||
|
||||
|
||||
def test_upload_falls_back_when_save_fails(
|
||||
client: TestClient, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
"""落盘失败时降级:仍 202 提交入库,但 metadata 不含落盘信息"""
|
||||
manager = FakeManager(task_id="fallback-task")
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
blocker = tmp_path / "blocker_file"
|
||||
blocker.write_bytes(b"x")
|
||||
monkeypatch.setattr(document_module.settings, "upload_dir", str(blocker))
|
||||
|
||||
resp = client.post(
|
||||
"/api/v1/documents/upload",
|
||||
files={"file": ("notes.txt", b"hello world", "text/plain")},
|
||||
)
|
||||
|
||||
assert resp.status_code == 202
|
||||
body = resp.json()
|
||||
assert body["data"]["saved_path"] == ""
|
||||
assert len(manager.submitted) == 1
|
||||
doc = manager.submitted[0]
|
||||
assert "raw_file_path" not in doc.metadata
|
||||
assert "original_filename" not in doc.metadata
|
||||
|
||||
|
||||
def test_upload_saved_path_uses_date_shard(
|
||||
client: TestClient, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""落盘路径使用 YYYY/MM 日期分片子目录"""
|
||||
manager = FakeManager()
|
||||
_inject_manager(monkeypatch, manager)
|
||||
|
||||
resp = client.post(
|
||||
"/api/v1/documents/upload",
|
||||
files={"file": ("shard.txt", b"shard content", "text/plain")},
|
||||
)
|
||||
|
||||
assert resp.status_code == 202
|
||||
saved_path = resp.json()["data"]["saved_path"]
|
||||
# 路径形如 .../uploads/YYYY/MM/{32位hex doc_id}_shard.txt
|
||||
norm = saved_path.replace("\\", "/")
|
||||
assert re.search(r"/\d{4}/\d{2}/[0-9a-f]{32}_shard\.txt$", norm), saved_path
|
||||
# 文件确已落盘到分片目录
|
||||
assert Path(saved_path).read_bytes() == b"shard content"
|
||||
@@ -0,0 +1,312 @@
|
||||
"""file_parser 单元测试:覆盖 txt/md/html/pdf/docx + 损坏文件 + 不支持扩展名 + PDF OCR 降级"""
|
||||
|
||||
import io
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from app.config import settings
|
||||
from app.core import file_parser as fp_module
|
||||
from app.core.file_parser import parse_file, supported_extensions
|
||||
|
||||
|
||||
def _make_minimal_pdf(text: str = "Hello PDF World") -> bytes:
|
||||
"""构造一个含一页文本的最小 PDF(pypdf 可读出文本)"""
|
||||
content_stream = f"BT /F1 24 Tf 100 700 Td ({text}) Tj ET".encode("latin-1")
|
||||
content_obj = b"<< /Length " + str(len(content_stream)).encode() + b" >>\nstream\n" + content_stream + b"\nendstream"
|
||||
return (
|
||||
b"%PDF-1.0\n"
|
||||
b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n"
|
||||
b"2 0 obj\n<< /Type /Pages /Kids [3 0 R] /Count 1 >>\nendobj\n"
|
||||
b"3 0 obj\n<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] "
|
||||
b"/Resources << /Font << /F1 4 0 R >> >> /Contents 5 0 R >>\nendobj\n"
|
||||
b"4 0 obj\n<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>\nendobj\n"
|
||||
b"5 0 obj\n" + content_obj + b"\nendobj\n"
|
||||
b"xref\n0 6\n"
|
||||
b"0000000000 65535 f\n"
|
||||
b"0000000010 00000 n\n"
|
||||
b"0000000059 00000 n\n"
|
||||
b"0000000115 00000 n\n"
|
||||
b"0000000241 00000 n\n"
|
||||
b"0000000316 00000 n\n"
|
||||
b"trailer\n<< /Size 6 /Root 1 0 R >>\nstartxref\n414\n%%EOF\n"
|
||||
)
|
||||
|
||||
|
||||
def _make_docx(text_lines: list[str]) -> bytes:
|
||||
from docx import Document # type: ignore[import-untyped]
|
||||
document = Document()
|
||||
for line in text_lines:
|
||||
document.add_paragraph(line)
|
||||
buf = io.BytesIO()
|
||||
document.save(buf)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def test_parse_txt_returns_decoded_text() -> None:
|
||||
"""txt:UTF-8 解码(含中文),无效字节 errors=replace 不抛错"""
|
||||
text = "你好世界 hello"
|
||||
assert parse_file("note.txt", text.encode("utf-8")) == text
|
||||
# 无效 UTF-8 字节不抛错(errors=replace 兜底)
|
||||
result = parse_file("bad.txt", b"\xff\xfe\x00invalid")
|
||||
assert isinstance(result, str)
|
||||
|
||||
|
||||
def test_parse_md_returns_decoded_text() -> None:
|
||||
"""md:与 txt 走同一解析器"""
|
||||
text = "# 标题\n\n正文内容"
|
||||
assert parse_file("note.md", text.encode("utf-8")) == text
|
||||
|
||||
|
||||
def test_parse_html_strips_tags() -> None:
|
||||
"""html:剥离 script/style 与所有标签,仅保留可见文本"""
|
||||
html = b"<html><body><h1>Hello</h1><p>World</p><script>x=1</script><style>p{}</style></body></html>"
|
||||
result = parse_file("page.html", html)
|
||||
assert "Hello" in result
|
||||
assert "World" in result
|
||||
assert "<" not in result
|
||||
assert ">" not in result
|
||||
assert "x=1" not in result
|
||||
assert "p{}" not in result
|
||||
|
||||
|
||||
def test_parse_htm_same_as_html() -> None:
|
||||
"""htm:与 html 走同一解析器"""
|
||||
html = b"<html><body><p>same content</p></body></html>"
|
||||
assert parse_file("page.htm", html) == parse_file("page.html", html)
|
||||
|
||||
|
||||
def test_parse_pdf_extracts_text() -> None:
|
||||
"""pdf:从最小 PDF 中提取文本"""
|
||||
pdf_bytes = _make_minimal_pdf("Hello PDF World")
|
||||
result = parse_file("doc.pdf", pdf_bytes)
|
||||
assert "Hello PDF World" in result
|
||||
|
||||
|
||||
def test_parse_docx_extracts_paragraphs() -> None:
|
||||
"""docx:提取段落文本"""
|
||||
docx_bytes = _make_docx(["第一段落", "第二段落"])
|
||||
result = parse_file("doc.docx", docx_bytes)
|
||||
assert "第一段落" in result
|
||||
assert "第二段落" in result
|
||||
|
||||
|
||||
def test_parse_file_unsupported_extension_raises() -> None:
|
||||
"""不支持的扩展名抛 ValueError,message 含扩展名"""
|
||||
with pytest.raises(ValueError, match=r"不支持的文件类型: \.xlsx"):
|
||||
parse_file("data.xlsx", b"binary")
|
||||
|
||||
|
||||
def test_parse_file_no_extension_raises() -> None:
|
||||
"""无扩展名抛 ValueError"""
|
||||
with pytest.raises(ValueError, match=r"不支持的文件类型"):
|
||||
parse_file("noext", b"text")
|
||||
|
||||
|
||||
def test_parse_file_corrupted_pdf_raises() -> None:
|
||||
"""损坏的 PDF 抛 ValueError,message 含'文件解析失败'"""
|
||||
with pytest.raises(ValueError, match=r"文件解析失败"):
|
||||
parse_file("bad.pdf", b"not a real pdf")
|
||||
|
||||
|
||||
def test_parse_file_empty_html_returns_empty_string() -> None:
|
||||
"""空 HTML 返回空字符串"""
|
||||
assert parse_file("empty.html", b"<html></html>") == ""
|
||||
|
||||
|
||||
def test_parse_file_html_with_script_style_excluded() -> None:
|
||||
"""HTML 中 script/style 内容被排除"""
|
||||
html = b"<html><body><p>visible</p><script>alert(1)</script><style>body{}</style></body></html>"
|
||||
result = parse_file("page.html", html)
|
||||
assert "visible" in result
|
||||
assert "alert" not in result
|
||||
assert "body{}" not in result
|
||||
|
||||
|
||||
def test_parse_file_html_with_entities_decoded() -> None:
|
||||
"""HTML 实体被解码"""
|
||||
html = b"<html><body><p>Tom & Jerry</p></body></html>"
|
||||
result = parse_file("page.html", html)
|
||||
assert "Tom & Jerry" in result
|
||||
assert "&" not in result
|
||||
|
||||
|
||||
def test_supported_extensions_contains_expected_set() -> None:
|
||||
"""supported_extensions 返回包含全部六种扩展名的集合"""
|
||||
exts = supported_extensions()
|
||||
assert {".txt", ".md", ".html", ".htm", ".pdf", ".docx"} <= exts
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# PDF OCR 降级路径测试(mock pypdf / pypdfium2 / rapidocr_onnxruntime,不真实下载模型)
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_ocr_state() -> Any:
|
||||
"""每个 OCR 测试前后重置模块级 OCR 引擎状态,避免相互污染"""
|
||||
saved_engine = fp_module._ocr_engine
|
||||
saved_unavailable = fp_module._ocr_unavailable
|
||||
yield
|
||||
fp_module._ocr_engine = saved_engine
|
||||
fp_module._ocr_unavailable = saved_unavailable
|
||||
|
||||
|
||||
class _FakeTextPage:
|
||||
"""pypdf PageObject 替身:返回固定文本"""
|
||||
|
||||
def __init__(self, text: str) -> None:
|
||||
self._text = text
|
||||
|
||||
def extract_text(self) -> str:
|
||||
return self._text
|
||||
|
||||
|
||||
class _FakePdfReader:
|
||||
"""pypdf.PdfReader 替身:构造时不解析,按预设页文本返回"""
|
||||
|
||||
def __init__(self, stream: Any) -> None:
|
||||
self.pages = [_FakeTextPage(""), _FakeTextPage("")]
|
||||
|
||||
|
||||
class _FakePilImage:
|
||||
pass
|
||||
|
||||
|
||||
class _FakeRenderResult:
|
||||
def to_pil(self) -> _FakePilImage:
|
||||
return _FakePilImage()
|
||||
|
||||
|
||||
class _FakePdfiumPage:
|
||||
def render(self, scale: float) -> _FakeRenderResult:
|
||||
return _FakeRenderResult()
|
||||
|
||||
|
||||
class _FakePdfDocument:
|
||||
"""pypdfium2.PdfDocument 替身"""
|
||||
|
||||
def __init__(self, stream: Any) -> None:
|
||||
self._n_pages = 2
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self._n_pages
|
||||
|
||||
def __getitem__(self, i: int) -> _FakePdfiumPage:
|
||||
return _FakePdfiumPage()
|
||||
|
||||
def close(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class _FakeOcrEngine:
|
||||
"""rapidocr RapidOCR 替身:每次返回固定识别结果"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.call_count = 0
|
||||
|
||||
def __call__(self, image: Any) -> tuple[list[list[Any]], float]:
|
||||
self.call_count += 1
|
||||
# 返回 [[box, text, score], ...] 结构
|
||||
return [[[0, 0], f"OCR文本第{self.call_count}页", 0.95]], 0.1
|
||||
|
||||
|
||||
def _patch_pdf_ocr_deps(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""把 _parse_pdf/_ocr_pdf 内部用到的 pypdf / pypdfium2 / rapidocr_onnxruntime 全部替换"""
|
||||
monkeypatch.setattr("pypdf.PdfReader", _FakePdfReader)
|
||||
monkeypatch.setattr("pypdfium2.PdfDocument", _FakePdfDocument)
|
||||
|
||||
fake_module = type(sys)("rapidocr_onnxruntime")
|
||||
fake_module.RapidOCR = _FakeOcrEngine
|
||||
monkeypatch.setitem(sys.modules, "rapidocr_onnxruntime", fake_module)
|
||||
|
||||
|
||||
def test_parse_pdf_ocr_fallback_when_text_layer_empty(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""扫描件 PDF(文本层全空)触发 OCR 降级,返回识别文本"""
|
||||
_patch_pdf_ocr_deps(monkeypatch)
|
||||
monkeypatch.setattr(settings, "pdf_ocr_enabled", True)
|
||||
monkeypatch.setattr(settings, "pdf_ocr_max_pages", 30)
|
||||
monkeypatch.setattr(settings, "pdf_ocr_dpi", 200)
|
||||
|
||||
result = parse_file("scan.pdf", b"fake pdf bytes")
|
||||
# 两页都跑了 OCR,每页返回一段文本
|
||||
assert "OCR文本第1页" in result
|
||||
assert "OCR文本第2页" in result
|
||||
|
||||
|
||||
def test_parse_pdf_ocr_skipped_when_disabled(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""pdf_ocr_enabled=False:文本层为空时直接返回空,不调 OCR"""
|
||||
_patch_pdf_ocr_deps(monkeypatch)
|
||||
monkeypatch.setattr(settings, "pdf_ocr_enabled", False)
|
||||
|
||||
result = parse_file("scan.pdf", b"fake pdf bytes")
|
||||
assert result == ""
|
||||
|
||||
|
||||
def test_parse_pdf_ocr_respects_max_pages(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""max_pages=1:只 OCR 第一页,第二页跳过"""
|
||||
_patch_pdf_ocr_deps(monkeypatch)
|
||||
monkeypatch.setattr(settings, "pdf_ocr_enabled", True)
|
||||
monkeypatch.setattr(settings, "pdf_ocr_max_pages", 1)
|
||||
monkeypatch.setattr(settings, "pdf_ocr_dpi", 200)
|
||||
|
||||
result = parse_file("scan.pdf", b"fake pdf bytes")
|
||||
assert "OCR文本第1页" in result
|
||||
assert "OCR文本第2页" not in result
|
||||
|
||||
|
||||
def test_parse_pdf_ocr_returns_empty_when_dependency_unavailable(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""rapidocr 导入失败:降级返回空文本,且把 _ocr_unavailable 置 True 避免重试"""
|
||||
monkeypatch.setattr("pypdf.PdfReader", _FakePdfReader)
|
||||
# 故意让 rapidocr_onnxruntime 提供一个非类的 RapidOCR,构造时抛错
|
||||
fake_module = type(sys)("rapidocr_onnxruntime")
|
||||
|
||||
def _boom(*args: Any, **kwargs: Any) -> None:
|
||||
raise RuntimeError("model missing")
|
||||
|
||||
fake_module.RapidOCR = _boom # type: ignore[attr-defined]
|
||||
monkeypatch.setitem(sys.modules, "rapidocr_onnxruntime", fake_module)
|
||||
monkeypatch.setattr(settings, "pdf_ocr_enabled", True)
|
||||
|
||||
assert parse_file("scan.pdf", b"fake pdf bytes") == ""
|
||||
assert fp_module._ocr_unavailable is True
|
||||
|
||||
|
||||
def test_parse_pdf_ocr_runtime_exception_falls_back_to_empty(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""OCR 运行时抛错:仅告警,降级返回空文本(不抛出 ValueError)"""
|
||||
monkeypatch.setattr("pypdf.PdfReader", _FakePdfReader)
|
||||
|
||||
class _ExplodingPdfDocument:
|
||||
def __init__(self, stream: Any) -> None:
|
||||
raise RuntimeError("pdfium render failed")
|
||||
|
||||
monkeypatch.setattr("pypdfium2.PdfDocument", _ExplodingPdfDocument)
|
||||
|
||||
fake_module = type(sys)("rapidocr_onnxruntime")
|
||||
fake_module.RapidOCR = _FakeOcrEngine
|
||||
monkeypatch.setitem(sys.modules, "rapidocr_onnxruntime", fake_module)
|
||||
monkeypatch.setattr(settings, "pdf_ocr_enabled", True)
|
||||
|
||||
assert parse_file("scan.pdf", b"fake pdf bytes") == ""
|
||||
|
||||
|
||||
def test_parse_pdf_text_layer_present_skips_ocr(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""文本层非空:直接返回文本,OCR 引擎不会被实例化"""
|
||||
call_count = 0
|
||||
|
||||
class _CountingReader:
|
||||
def __init__(self, stream: Any) -> None:
|
||||
self.pages = [_FakeTextPage("这是文本层的内容")]
|
||||
|
||||
monkeypatch.setattr("pypdf.PdfReader", _CountingReader)
|
||||
|
||||
# 即便 OCR 依赖故意坏掉,也不应被调用
|
||||
bad_module = type(sys)("rapidocr_onnxruntime")
|
||||
bad_module.RapidOCR = lambda *a, **k: (_ for _ in ()).throw(RuntimeError("不应被调用")) # type: ignore[attr-defined]
|
||||
monkeypatch.setitem(sys.modules, "rapidocr_onnxruntime", bad_module)
|
||||
monkeypatch.setattr(settings, "pdf_ocr_enabled", True)
|
||||
|
||||
result = parse_file("text.pdf", b"fake pdf bytes")
|
||||
assert result == "这是文本层的内容"
|
||||
assert fp_module._ocr_engine is None
|
||||
@@ -1,12 +1,13 @@
|
||||
"""IngestTaskManager 单元测试(FakeIngester/FakeRedis,不真实联网)"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from app.config import Settings
|
||||
from app.core.ingest_tasks import REDIS_KEY_PREFIX, IngestTaskManager, IngestTaskStatus
|
||||
from app.core.ingest_tasks import DEDUP_KEY_PREFIX, REDIS_KEY_PREFIX, IngestTaskManager, IngestTaskStatus
|
||||
from app.core.ingestion import IngestionError
|
||||
from app.models.document import DocumentInput, DocumentSummary, IngestionResult, SummaryLevel
|
||||
|
||||
@@ -198,3 +199,91 @@ async def test_get_falls_back_to_redis_then_none() -> None:
|
||||
record = await manager.get("abc")
|
||||
assert record is not None
|
||||
assert record["status"] == "done"
|
||||
|
||||
|
||||
def _text_hash(text: str) -> str:
|
||||
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
async def test_dedup_hit_reuses_old_doc_id_and_skips_pipeline() -> None:
|
||||
"""去重命中:相同 text 第二次提交直接 done,复用旧 doc_id,deduplicated=True,不调 ingester"""
|
||||
ingester = FakeIngester()
|
||||
redis = FakeRedis()
|
||||
manager = IngestTaskManager(ingester, redis, Settings())
|
||||
|
||||
# 首次提交:正常跑流水线
|
||||
text = "重复内容"
|
||||
task1 = await manager.submit(DocumentInput(text=text, title="t1"))
|
||||
final1 = await manager.wait_done(task1, timeout=5)
|
||||
assert final1["status"] == IngestTaskStatus.DONE
|
||||
assert final1["result"]["document_id"] == "doc-1"
|
||||
assert final1["result"]["deduplicated"] is False
|
||||
assert len(ingester.calls) == 1
|
||||
|
||||
# dedup 记录已写入 Redis
|
||||
stored = await redis.get_json(f"{DEDUP_KEY_PREFIX}{_text_hash(text)}")
|
||||
assert stored is not None
|
||||
assert stored["document_id"] == "doc-1"
|
||||
|
||||
# 第二次提交相同 text:直接 done,不重跑流水线
|
||||
task2 = await manager.submit(DocumentInput(text=text, title="t2"))
|
||||
# wait_done 会先 await 所有 fire-and-forget 镜像任务再返回,确保 Redis 已写
|
||||
record2 = await manager.wait_done(task2, timeout=5)
|
||||
assert record2["status"] == IngestTaskStatus.DONE
|
||||
assert record2["result"]["document_id"] == "doc-1"
|
||||
assert record2["result"]["deduplicated"] is True
|
||||
# ingester 没被再次调用
|
||||
assert len(ingester.calls) == 1
|
||||
# 镜像已写入 Redis
|
||||
mirrored = await redis.get_json(f"{REDIS_KEY_PREFIX}{task2}")
|
||||
assert mirrored is not None
|
||||
assert mirrored["status"] == IngestTaskStatus.DONE
|
||||
assert mirrored["result"]["deduplicated"] is True
|
||||
|
||||
|
||||
async def test_dedup_miss_when_text_differs() -> None:
|
||||
"""去重未命中:不同 text 走原 _run,完成后写入对应 dedup key"""
|
||||
ingester = FakeIngester()
|
||||
redis = FakeRedis()
|
||||
manager = IngestTaskManager(ingester, redis, Settings())
|
||||
|
||||
task1 = await manager.submit(DocumentInput(text="内容A", title="t1"))
|
||||
await manager.wait_done(task1, timeout=5)
|
||||
task2 = await manager.submit(DocumentInput(text="内容B", title="t2"))
|
||||
await manager.wait_done(task2, timeout=5)
|
||||
|
||||
assert await redis.get_json(f"{DEDUP_KEY_PREFIX}{_text_hash('内容A')}") is not None
|
||||
assert await redis.get_json(f"{DEDUP_KEY_PREFIX}{_text_hash('内容B')}") is not None
|
||||
assert len(ingester.calls) == 2
|
||||
|
||||
|
||||
async def test_dedup_skipped_when_redis_unavailable() -> None:
|
||||
"""Redis 不可用:跳过去重,相同 text 仍走完整流水线"""
|
||||
ingester = FakeIngester()
|
||||
manager = IngestTaskManager(ingester, None, Settings())
|
||||
|
||||
task1 = await manager.submit(DocumentInput(text="内容", title="t1"))
|
||||
final1 = await manager.wait_done(task1, timeout=5)
|
||||
assert final1["result"]["deduplicated"] is False
|
||||
|
||||
task2 = await manager.submit(DocumentInput(text="内容", title="t2"))
|
||||
final2 = await manager.wait_done(task2, timeout=5)
|
||||
assert final2["status"] == IngestTaskStatus.DONE
|
||||
assert final2["result"]["deduplicated"] is False
|
||||
assert len(ingester.calls) == 2
|
||||
|
||||
|
||||
async def test_dedup_lookup_failure_falls_back_to_normal_pipeline() -> None:
|
||||
"""Redis get_json 抛错:去重查询降级为未命中,走原流水线"""
|
||||
class _ExplodingRedis(FakeRedis):
|
||||
async def get_json(self, key: str) -> dict[str, Any] | None:
|
||||
raise RuntimeError("redis down")
|
||||
|
||||
ingester = FakeIngester()
|
||||
redis = _ExplodingRedis()
|
||||
manager = IngestTaskManager(ingester, redis, Settings())
|
||||
|
||||
task = await manager.submit(DocumentInput(text="x", title="t"))
|
||||
final = await manager.wait_done(task, timeout=5)
|
||||
assert final["status"] == IngestTaskStatus.DONE
|
||||
assert len(ingester.calls) == 1
|
||||
|
||||
@@ -148,6 +148,69 @@ class TestQueryParserParse:
|
||||
|
||||
assert [c.name for c in parsed.categories] == ["财务行政"]
|
||||
|
||||
async def test_parse_new_fields_populated(self):
|
||||
"""JSON 含 entities/intent/time_range → 正确填充字段"""
|
||||
response = json.dumps(
|
||||
{
|
||||
"categories": [{"name": "技术文档", "confidence": 0.9}],
|
||||
"rewrite": "如何设计微服务架构",
|
||||
"keywords": ["架构"],
|
||||
"entities": ["微服务", "Kubernetes"],
|
||||
"intent": "操作",
|
||||
"time_range": "2023年",
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
parser = _make_parser(response)
|
||||
|
||||
parsed = await parser.parse("怎么做微服务部署")
|
||||
|
||||
assert parsed.parse_failed is False
|
||||
assert parsed.entities == ["微服务", "Kubernetes"]
|
||||
assert parsed.intent == "操作"
|
||||
assert parsed.time_range == "2023年"
|
||||
|
||||
async def test_parse_new_fields_missing_defaults(self):
|
||||
"""JSON 缺 entities/intent/time_range → 降级为默认值,parse_failed=False(向后兼容)"""
|
||||
response = json.dumps(
|
||||
{
|
||||
"categories": [{"name": "技术文档", "confidence": 0.9}],
|
||||
"rewrite": "如何设计架构",
|
||||
"keywords": ["架构"],
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
parser = _make_parser(response)
|
||||
|
||||
parsed = await parser.parse("怎么做架构设计")
|
||||
|
||||
assert parsed.parse_failed is False
|
||||
assert parsed.entities == []
|
||||
assert parsed.intent == ""
|
||||
assert parsed.time_range == ""
|
||||
|
||||
async def test_parse_new_fields_wrong_type_defaults(self):
|
||||
"""entities 非列表 / intent 非字符串 / time_range 非字符串 → 置空,parse_failed=False(容错)"""
|
||||
response = json.dumps(
|
||||
{
|
||||
"categories": [{"name": "技术文档", "confidence": 0.9}],
|
||||
"rewrite": "如何设计架构",
|
||||
"keywords": ["架构"],
|
||||
"entities": "微服务", # 非列表
|
||||
"intent": 123, # 非字符串
|
||||
"time_range": ["2023"], # 非字符串
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
parser = _make_parser(response)
|
||||
|
||||
parsed = await parser.parse("怎么做架构设计")
|
||||
|
||||
assert parsed.parse_failed is False
|
||||
assert parsed.entities == []
|
||||
assert parsed.intent == ""
|
||||
assert parsed.time_range == ""
|
||||
|
||||
|
||||
class TestDecideRoute:
|
||||
"""decide_route 纯函数的四个分支"""
|
||||
|
||||
@@ -112,9 +112,9 @@ class _CountingRetriever:
|
||||
return self.response
|
||||
|
||||
|
||||
def _search_cache_key(query: str, top_k: int | None = None) -> str:
|
||||
"""与路由侧一致的检索缓存键"""
|
||||
return f"search:{sha256((query + '|' + str(top_k)).encode()).hexdigest()[:16]}"
|
||||
def _search_cache_key(query: str, top_k: int | None = None, summarize: bool = False) -> str:
|
||||
"""与路由侧一致的检索缓存键(query + top_k + summarize 三者均参与键计算)"""
|
||||
return f"search:{sha256((query + '|' + str(top_k) + '|' + str(summarize)).encode()).hexdigest()[:16]}"
|
||||
|
||||
|
||||
class TestSearchApiCache:
|
||||
|
||||
@@ -0,0 +1,138 @@
|
||||
"""ResultSummarizer 单元测试(FakeOllama,不依赖真实 Ollama)
|
||||
|
||||
覆盖:
|
||||
- summarize 正常:调用 ollama,prompt 含 query 与 hits 文本,返回 strip 后的 summary
|
||||
- summarize 无 hits:返回空串且不调 ollama
|
||||
- summarize ollama 异常:返回空串
|
||||
- _build_context 拼接格式:编号 / title / section_path / 文本
|
||||
- 取前 settings.result_summary_max_hits 条:超量 hits 只取前 N 条进入 prompt
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from app.config import settings
|
||||
from app.core.result_summarizer import ResultSummarizer
|
||||
from app.models.search import SearchHit
|
||||
|
||||
|
||||
class FakeOllama:
|
||||
"""记录调用并返回固定响应的假 OllamaClient"""
|
||||
|
||||
def __init__(self, response: str = "这是总结") -> None:
|
||||
self.response = response
|
||||
self.calls: list[dict] = []
|
||||
|
||||
async def generate(self, prompt: str, json_mode: bool = False) -> str:
|
||||
self.calls.append({"prompt": prompt, "json_mode": json_mode})
|
||||
return self.response
|
||||
|
||||
|
||||
class FailingOllama:
|
||||
"""generate 抛异常的假 OllamaClient,用于测试容错降级"""
|
||||
|
||||
async def generate(self, prompt: str, json_mode: bool = False) -> str:
|
||||
raise RuntimeError("ollama 不可用")
|
||||
|
||||
|
||||
def _hit(idx: int, title: str = "", section_path: str = "") -> SearchHit:
|
||||
"""构造测试用 SearchHit"""
|
||||
return SearchHit(
|
||||
text=f"文本内容-{idx}",
|
||||
doc_id=f"doc-{idx}",
|
||||
title=title or f"标题-{idx}",
|
||||
section_path=section_path,
|
||||
score=0.1 * idx,
|
||||
)
|
||||
|
||||
|
||||
class TestSummarize:
|
||||
"""ResultSummarizer.summarize 主流程"""
|
||||
|
||||
async def test_summarize_normal(self):
|
||||
"""正常总结:调用 ollama,prompt 含 query 与 hits 文本,返回 strip 后的 summary"""
|
||||
ollama = FakeOllama(response=" 这是 AI 总结 ")
|
||||
summarizer = ResultSummarizer(ollama=ollama) # type: ignore[arg-type]
|
||||
hits = [_hit(1, "文档A", "章节A"), _hit(2, "文档B", "章节B")]
|
||||
|
||||
result = await summarizer.summarize("安装步骤是什么", hits)
|
||||
|
||||
assert result == "这是 AI 总结"
|
||||
assert len(ollama.calls) == 1
|
||||
prompt = ollama.calls[0]["prompt"]
|
||||
assert "安装步骤是什么" in prompt
|
||||
assert "文本内容-1" in prompt
|
||||
assert "文本内容-2" in prompt
|
||||
assert "文档A" in prompt
|
||||
assert "文档B" in prompt
|
||||
assert "章节A" in prompt
|
||||
assert "章节B" in prompt
|
||||
|
||||
async def test_summarize_empty_hits_returns_empty(self):
|
||||
"""无 hits:返回空串且不调用 ollama"""
|
||||
ollama = FakeOllama()
|
||||
summarizer = ResultSummarizer(ollama=ollama) # type: ignore[arg-type]
|
||||
|
||||
result = await summarizer.summarize("任何问题", [])
|
||||
|
||||
assert result == ""
|
||||
assert ollama.calls == []
|
||||
|
||||
async def test_summarize_ollama_error_returns_empty(self):
|
||||
"""ollama.generate 抛异常:返回空串(容错降级)"""
|
||||
summarizer = ResultSummarizer(ollama=FailingOllama()) # type: ignore[arg-type]
|
||||
hits = [_hit(1)]
|
||||
|
||||
result = await summarizer.summarize("问题", hits)
|
||||
|
||||
assert result == ""
|
||||
|
||||
|
||||
class TestBuildContext:
|
||||
"""_build_context 拼接格式:编号 / title / section_path / 文本"""
|
||||
|
||||
def test_build_context_format(self):
|
||||
hits = [
|
||||
_hit(1, "文档A", "第一章"),
|
||||
_hit(2, "文档B", "第二章"),
|
||||
]
|
||||
context = ResultSummarizer._build_context(hits)
|
||||
|
||||
# 每条带编号、title、section_path
|
||||
assert "[1] / 文档A / 第一章" in context
|
||||
assert "[2] / 文档B / 第二章" in context
|
||||
# 文本内容拼接
|
||||
assert "文本内容-1" in context
|
||||
assert "文本内容-2" in context
|
||||
# 分隔符
|
||||
assert "---" in context
|
||||
|
||||
def test_build_context_empty_title_and_section(self):
|
||||
"""title 与 section_path 为空时头部仅保留编号"""
|
||||
hit = SearchHit(text="纯文本", doc_id="d1", title="", section_path="", score=1.0)
|
||||
context = ResultSummarizer._build_context([hit])
|
||||
assert context.strip().startswith("[1]")
|
||||
assert "纯文本" in context
|
||||
|
||||
|
||||
class TestMaxHitsLimit:
|
||||
"""取前 settings.result_summary_max_hits 条命中"""
|
||||
|
||||
async def test_only_first_n_hits_in_prompt(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""构造超过 max_hits 的 hits,验证 prompt 只含前 N 条文本"""
|
||||
max_hits = 3
|
||||
monkeypatch.setattr(settings, "result_summary_max_hits", max_hits)
|
||||
ollama = FakeOllama(response="总结")
|
||||
summarizer = ResultSummarizer(ollama=ollama) # type: ignore[arg-type]
|
||||
|
||||
# 构造 max_hits + 5 条 hits
|
||||
hits = [_hit(i) for i in range(max_hits + 5)]
|
||||
await summarizer.summarize("问题", hits)
|
||||
|
||||
assert len(ollama.calls) == 1
|
||||
prompt = ollama.calls[0]["prompt"]
|
||||
# 前 max_hits 条文本出现在 prompt 中
|
||||
for i in range(max_hits):
|
||||
assert f"文本内容-{i}" in prompt
|
||||
# 超出的文本不出现在 prompt 中
|
||||
for i in range(max_hits, max_hits + 5):
|
||||
assert f"文本内容-{i}" not in prompt
|
||||
Reference in New Issue
Block a user