feat: 新增多格式文件上传入库与认证体系

- 新增 JWT 认证模块,支持登录/注册/用户管理
- 新增文件上传接口,支持 .txt/.md/.html/.pdf/.docx 等格式解析入库
- 新增检索结果 AI 总结功能
- 新增文本去重缓存机制
- 新增全局认证夹具简化测试
- 新增配置项与环境变量支持
- 完善文档与测试覆盖
This commit is contained in:
2026-07-30 10:30:15 +08:00
parent 51dc8dc4f6
commit dce9e31bde
31 changed files with 3021 additions and 45 deletions
+27
View File
@@ -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 掉对应 overrideautouse 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()
+293
View File
@@ -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_adminadmin 通过、非 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
+313
View File
@@ -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=1001message 含'大小上限',未提交任务"""
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=1001message 含'无法从文件提取文本',未提交任务"""
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:
"""损坏的 PDFcode=1001message 含'文件解析失败',未提交任务"""
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"
+312
View File
@@ -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:
"""不支持的扩展名抛 ValueErrormessage 含扩展名"""
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 抛 ValueErrormessage 含'文件解析失败'"""
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 &amp; Jerry</p></body></html>"
result = parse_file("page.html", html)
assert "Tom & Jerry" in result
assert "&amp;" 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
+90 -1
View File
@@ -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_iddeduplicated=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
+63
View File
@@ -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 纯函数的四个分支"""
+3 -3
View File
@@ -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:
+138
View File
@@ -0,0 +1,138 @@
"""ResultSummarizer 单元测试(FakeOllama,不依赖真实 Ollama
覆盖:
- summarize 正常:调用 ollamaprompt 含 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):
"""正常总结:调用 ollamaprompt 含 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