kplam
2ab8b56a01
此提交实现了完整的知识库管理系统： 1. 新增Vue3 + Antd Vue前端管理后台，包含登录、文档管理、检索、类目设置等完整页面 2. 重构后端LLM调用抽象层，支持Ollama与OpenAI兼容服务动态切换 3. 调整默认嵌入模型配置为本地bge-m3模式 4. 优化入库任务去重逻辑与缓存清理机制 5. 完善Docker镜像构建与docker-compose部署配置 6. 修复多项测试用例与兼容性问题 7. 新增运行时配置API，支持动态调整系统参数
346 lines
12 KiB
Python
346 lines
12 KiB
Python
"""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_returns_empty(
|
||
monkeypatch: pytest.MonkeyPatch,
|
||
) -> None:
|
||
"""损坏的 PDF：插件化后文本层提取失败被捕获，OCR 关闭时返回空字符串
|
||
|
||
原 test_parse_file_corrupted_pdf_raises 期望 ValueError，但插件化重构后
|
||
file_parser 设计为优雅降级（pypdf 失败 → OCR 兜底 → 都失败返回空），
|
||
不再向上抛异常。关闭 OCR 避免触发 rapidocr 模型下载拖慢测试。
|
||
"""
|
||
monkeypatch.setattr(settings, "pdf_ocr_enabled", False)
|
||
assert 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 测试前后重置 RapidocrOcrEngine/TesseractOcrEngine 类级状态与插件缓存
|
||
|
||
file_parser 插件化后，模块级 _ocr_engine/_ocr_unavailable 已移除，
|
||
RapidocrOcrEngine 用类级字段 _engine/_unavailable 单例化。
|
||
测试前重置为干净状态（避免上个测试残留），测试后清理插件缓存。
|
||
"""
|
||
# 测试前：重置为干净状态
|
||
fp_module.RapidocrOcrEngine._engine = None
|
||
fp_module.RapidocrOcrEngine._unavailable = False
|
||
fp_module.TesseractOcrEngine._unavailable = False
|
||
fp_module._ocr_plugin_cache.clear()
|
||
yield
|
||
# 测试后：再次清理，避免污染后续非 OCR 测试
|
||
fp_module.RapidocrOcrEngine._engine = None
|
||
fp_module.RapidocrOcrEngine._unavailable = False
|
||
fp_module.TesseractOcrEngine._unavailable = False
|
||
fp_module._ocr_plugin_cache.clear()
|
||
|
||
|
||
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.RapidocrOcrEngine._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.RapidocrOcrEngine._engine is None
|