chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,322 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import time
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import fitz
|
||||
import pandas as pd
|
||||
import pytest
|
||||
import yuxi.knowledge.parser.unified as parser_unified
|
||||
from docx import Document
|
||||
from PIL import Image
|
||||
|
||||
from yuxi.knowledge.parser import Parser
|
||||
from yuxi.knowledge.parser.factory import DocumentProcessorFactory
|
||||
|
||||
DATA_DIR = Path(__file__).resolve().parents[2] / "data"
|
||||
|
||||
|
||||
def _build_pdf(file_path: Path, text: str) -> None:
|
||||
doc = fitz.open()
|
||||
page = doc.new_page()
|
||||
page.insert_text((72, 72), text)
|
||||
doc.save(str(file_path))
|
||||
doc.close()
|
||||
|
||||
|
||||
def _build_docx(file_path: Path, text: str) -> None:
|
||||
document = Document()
|
||||
document.add_paragraph(text)
|
||||
document.save(str(file_path))
|
||||
|
||||
|
||||
def _build_png(file_path: Path) -> None:
|
||||
image = Image.new("RGB", (120, 80), "white")
|
||||
image.save(str(file_path))
|
||||
|
||||
|
||||
def test_parser_parse_pdf_file_returns_markdown_text(tmp_path: Path):
|
||||
file_path = tmp_path / "parser_test.pdf"
|
||||
_build_pdf(file_path, "Parser PDF content")
|
||||
|
||||
markdown = Parser.parse(str(file_path), params={"ocr_engine": "disable"})
|
||||
|
||||
assert isinstance(markdown, str)
|
||||
assert "Parser" in markdown
|
||||
assert "content" in markdown
|
||||
assert len(markdown.strip()) > 0
|
||||
|
||||
|
||||
def test_parser_parse_docx_file_returns_markdown_text(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
file_path = tmp_path / "parser_test.docx"
|
||||
_build_docx(file_path, "Parser DOCX content")
|
||||
|
||||
# 避免测试依赖 docling 行为，直接验证统一 parser 可回退到 python-docx。
|
||||
def _raise_docling_error(*args, **kwargs):
|
||||
raise RuntimeError("force fallback to python-docx")
|
||||
|
||||
monkeypatch.setattr(parser_unified, "_convert_with_docling", _raise_docling_error)
|
||||
|
||||
markdown = Parser.parse(str(file_path))
|
||||
|
||||
assert isinstance(markdown, str)
|
||||
assert "Parser DOCX content" in markdown
|
||||
assert len(markdown.strip()) > 0
|
||||
|
||||
|
||||
def test_convert_csv_to_markdown_preserves_column_dtypes(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
file_path = tmp_path / "parser_test.csv"
|
||||
file_path.write_text("id,score\n9007199254740993,2.5\n", encoding="utf-8")
|
||||
captured_dtypes: list[dict[str, object]] = []
|
||||
original_to_markdown = pd.DataFrame.to_markdown
|
||||
|
||||
def _capture_dtypes(dataframe: pd.DataFrame, *args, **kwargs) -> str:
|
||||
captured_dtypes.append(dataframe.dtypes.to_dict())
|
||||
return original_to_markdown(dataframe, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(pd.DataFrame, "to_markdown", _capture_dtypes)
|
||||
|
||||
markdown = parser_unified._convert_csv_to_markdown(file_path)
|
||||
|
||||
assert markdown
|
||||
assert str(captured_dtypes[0]["id"]) == "int64"
|
||||
|
||||
|
||||
def test_convert_with_docling_reinserts_image_links_in_document_order(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
file_path = tmp_path / "parser_test.docx"
|
||||
file_path.write_bytes(b"fake docx")
|
||||
first_image = base64.b64encode(b"first image").decode()
|
||||
second_image = base64.b64encode(b"second image").decode()
|
||||
fake_doc = SimpleNamespace(
|
||||
pictures=[
|
||||
SimpleNamespace(image=SimpleNamespace(uri=f"data:image/png;base64,{first_image}")),
|
||||
SimpleNamespace(image=SimpleNamespace(uri="https://example.test/remote.png")),
|
||||
SimpleNamespace(image=SimpleNamespace(uri=f"data:image/png;base64,{second_image}")),
|
||||
],
|
||||
export_to_markdown=lambda: "before\n<!-- image -->\nremote\n<!-- image -->\nbetween\n<!-- image -->\nafter",
|
||||
)
|
||||
fake_result = SimpleNamespace(status=SimpleNamespace(name="SUCCESS"), document=fake_doc)
|
||||
uploaded_images: list[bytes] = []
|
||||
|
||||
class FakeConverter:
|
||||
def convert(self, path: Path):
|
||||
assert path == file_path
|
||||
return fake_result
|
||||
|
||||
def _fake_upload_image_to_minio(image_data, filename, bucket_name, object_prefix):
|
||||
uploaded_images.append(image_data)
|
||||
return f"https://example.test/{len(uploaded_images)}.png"
|
||||
|
||||
monkeypatch.setattr(parser_unified, "_get_docling_converter", lambda: FakeConverter())
|
||||
monkeypatch.setattr(parser_unified, "_upload_image_to_minio", _fake_upload_image_to_minio)
|
||||
image_timestamps = iter([1.0, 2.0])
|
||||
monkeypatch.setattr(parser_unified.time, "time", lambda: next(image_timestamps))
|
||||
|
||||
markdown = parser_unified._convert_with_docling(file_path)
|
||||
|
||||
assert uploaded_images == [b"first image", b"second image"]
|
||||
assert markdown == (
|
||||
"before\n"
|
||||
"\n"
|
||||
"remote\n"
|
||||
"\n"
|
||||
"between\n"
|
||||
"\n"
|
||||
"after"
|
||||
)
|
||||
|
||||
|
||||
def test_convert_with_docling_keeps_image_placeholder_when_upload_fails(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
file_path = tmp_path / "parser_test.docx"
|
||||
file_path.write_bytes(b"fake docx")
|
||||
image = base64.b64encode(b"image data").decode()
|
||||
fake_doc = SimpleNamespace(
|
||||
pictures=[SimpleNamespace(image=SimpleNamespace(uri=f"data:image/png;base64,{image}"))],
|
||||
export_to_markdown=lambda: "before\n<!-- image -->\nafter",
|
||||
)
|
||||
fake_result = SimpleNamespace(status=SimpleNamespace(name="SUCCESS"), document=fake_doc)
|
||||
|
||||
class FakeConverter:
|
||||
def convert(self, path: Path):
|
||||
assert path == file_path
|
||||
return fake_result
|
||||
|
||||
def _raise_upload_error(*args, **kwargs):
|
||||
raise RuntimeError("upload failed")
|
||||
|
||||
monkeypatch.setattr(parser_unified, "_get_docling_converter", lambda: FakeConverter())
|
||||
monkeypatch.setattr(parser_unified, "_upload_image_to_minio", _raise_upload_error)
|
||||
monkeypatch.setattr(parser_unified.time, "time", lambda: 1.0)
|
||||
|
||||
markdown = parser_unified._convert_with_docling(file_path)
|
||||
|
||||
assert markdown == "before\n[图片: image_1000000.png]\nafter"
|
||||
|
||||
|
||||
def test_parser_parse_png_file_returns_markdown_text_with_mocked_ocr(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
file_path = tmp_path / "parser_test.png"
|
||||
_build_png(file_path)
|
||||
|
||||
async def _fake_parse_image_async(file, params=None):
|
||||
return "Parser PNG content"
|
||||
|
||||
monkeypatch.setattr(parser_unified, "parse_image_async", _fake_parse_image_async)
|
||||
|
||||
markdown = Parser.parse(str(file_path), params={"ocr_engine": "rapid_ocr"})
|
||||
|
||||
assert isinstance(markdown, str)
|
||||
assert "Parser PNG content" in markdown
|
||||
assert len(markdown.strip()) > 0
|
||||
|
||||
|
||||
def test_parse_image_uses_ocr_engine_config(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
file_path = tmp_path / "parser_test.png"
|
||||
_build_png(file_path)
|
||||
captured = {}
|
||||
|
||||
def _fake_process_file(processor_type, file, params=None):
|
||||
captured["processor_type"] = processor_type
|
||||
captured["file"] = file
|
||||
captured["params"] = params
|
||||
return "OCR content"
|
||||
|
||||
monkeypatch.setattr(DocumentProcessorFactory, "process_file", _fake_process_file)
|
||||
|
||||
result = parser_unified.parse_image(
|
||||
str(file_path),
|
||||
params={
|
||||
"ocr_engine": "mineru_ocr",
|
||||
"backend": "old-backend",
|
||||
"ocr_engine_config": {"backend": "pipeline", "formula_enable": False},
|
||||
},
|
||||
)
|
||||
|
||||
assert result == "OCR content"
|
||||
assert captured["processor_type"] == "mineru_ocr"
|
||||
assert captured["file"] == str(file_path)
|
||||
assert captured["params"]["backend"] == "pipeline"
|
||||
assert captured["params"]["formula_enable"] is False
|
||||
|
||||
|
||||
def test_parse_image_ignores_enable_ocr(tmp_path: Path) -> None:
|
||||
file_path = tmp_path / "parser_test.png"
|
||||
_build_png(file_path)
|
||||
|
||||
with pytest.raises(ValueError, match="必须启用OCR"):
|
||||
parser_unified.parse_image(str(file_path), params={"ocr_engine": "disable", "enable_ocr": "rapid_ocr"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parser_aparse_pdf_file_returns_markdown_text(tmp_path: Path):
|
||||
file_path = tmp_path / "parser_test_async.pdf"
|
||||
_build_pdf(file_path, "Async Parser PDF content")
|
||||
|
||||
markdown = await Parser.aparse(str(file_path), params={"ocr_engine": "disable"})
|
||||
|
||||
assert isinstance(markdown, str)
|
||||
assert "Async" in markdown
|
||||
assert "content" in markdown
|
||||
assert len(markdown.strip()) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parser_aparse_docx_does_not_block_event_loop(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
file_path = tmp_path / "parser_test_async.docx"
|
||||
file_path.write_bytes(b"fake docx")
|
||||
completion_order: list[str] = []
|
||||
|
||||
def _slow_docling_conversion(*args, **kwargs) -> str:
|
||||
time.sleep(0.1)
|
||||
return "Async DOCX content"
|
||||
|
||||
async def _parse_document() -> None:
|
||||
await Parser.aparse(str(file_path))
|
||||
completion_order.append("parse")
|
||||
|
||||
async def _record_event_loop_progress() -> None:
|
||||
await asyncio.sleep(0.01)
|
||||
completion_order.append("event_loop")
|
||||
|
||||
monkeypatch.setattr(parser_unified, "_convert_with_docling", _slow_docling_conversion)
|
||||
|
||||
await asyncio.gather(_parse_document(), _record_event_loop_progress())
|
||||
|
||||
assert completion_order == ["event_loop", "parse"]
|
||||
|
||||
|
||||
def test_parse_pdf_uses_config_default_ocr_when_engine_missing(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import yuxi
|
||||
|
||||
file_path = tmp_path / "parser_test.pdf"
|
||||
_build_pdf(file_path, "Parser PDF content")
|
||||
captured = {}
|
||||
|
||||
def _fake_process_file(processor_type, file, params=None):
|
||||
captured["processor_type"] = processor_type
|
||||
captured["file"] = file
|
||||
captured["params"] = params
|
||||
return "default OCR content"
|
||||
|
||||
monkeypatch.setattr(yuxi.config, "default_ocr_engine", "mineru_ocr")
|
||||
monkeypatch.setattr(DocumentProcessorFactory, "process_file", _fake_process_file)
|
||||
|
||||
result = parser_unified.parse_pdf(str(file_path), params={})
|
||||
|
||||
assert result == "default OCR content"
|
||||
assert captured["processor_type"] == "mineru_ocr"
|
||||
assert captured["file"] == str(file_path)
|
||||
|
||||
|
||||
def test_parse_pdf_keeps_explicit_disable_when_default_ocr_enabled(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import yuxi
|
||||
|
||||
file_path = tmp_path / "parser_test.pdf"
|
||||
_build_pdf(file_path, "Parser PDF content")
|
||||
monkeypatch.setattr(yuxi.config, "default_ocr_engine", "mineru_ocr")
|
||||
|
||||
result = parser_unified.parse_pdf(str(file_path), params={"ocr_engine": "disable"})
|
||||
|
||||
assert "Parser PDF content" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parser_aparse_image_file_with_mineru_when_available():
|
||||
file_path = DATA_DIR / "测试图片.png"
|
||||
assert file_path.exists(), f"测试文件不存在: {file_path}"
|
||||
|
||||
health = await asyncio.to_thread(DocumentProcessorFactory.check_health, "mineru_ocr")
|
||||
if health.get("status") != "healthy":
|
||||
pytest.skip(f"mineru_ocr 不可用: {health.get('message', 'unknown')}")
|
||||
|
||||
markdown = await Parser.aparse(
|
||||
str(file_path),
|
||||
params={"ocr_engine": "mineru_ocr", "backend": "pipeline"},
|
||||
)
|
||||
|
||||
assert isinstance(markdown, str)
|
||||
assert len(markdown) > 100
|
||||
assert len(markdown.strip()) > 0
|
||||
Reference in New Issue
Block a user