kplam
dce9e31bde
- 新增 JWT 认证模块，支持登录/注册/用户管理 - 新增文件上传接口，支持 .txt/.md/.html/.pdf/.docx 等格式解析入库 - 新增检索结果 AI 总结功能 - 新增文本去重缓存机制 - 新增全局认证夹具简化测试 - 新增配置项与环境变量支持 - 完善文档与测试覆盖
314 lines
10 KiB
Python
314 lines
10 KiB
Python
"""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"
|