105 lines
3.4 KiB
Python
105 lines
3.4 KiB
Python
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from services.pdf_ocr_service import extract_scanned_pdf_text
|
|
|
|
|
|
class FakePixmap:
|
|
def tobytes(self, *_args, **_kwargs):
|
|
return b"jpeg-page"
|
|
|
|
|
|
class FakePage:
|
|
def get_pixmap(self, **_kwargs):
|
|
return FakePixmap()
|
|
|
|
|
|
class FakeDocument:
|
|
page_count = 2
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
return None
|
|
|
|
def load_page(self, _index):
|
|
return FakePage()
|
|
|
|
|
|
def test_scanned_pdf_ocr_preserves_page_order_and_reports_progress(monkeypatch):
|
|
fake_pymupdf = SimpleNamespace(
|
|
open=lambda _path: FakeDocument(),
|
|
Matrix=lambda x, y: (x, y),
|
|
csRGB="rgb",
|
|
)
|
|
monkeypatch.setitem(__import__("sys").modules, "pymupdf", fake_pymupdf)
|
|
progress = []
|
|
reservation = SimpleNamespace()
|
|
config = SimpleNamespace(
|
|
api_key="configured",
|
|
ocr_model="qwen-vl-ocr",
|
|
vision_model="vision",
|
|
vision_max_tokens=2048,
|
|
)
|
|
|
|
with (
|
|
patch("services.pdf_ocr_service.get_chat_model_config", return_value=config),
|
|
patch("services.pdf_ocr_service.prepare_image", return_value=SimpleNamespace()),
|
|
patch(
|
|
"services.pdf_ocr_service.call_vision_model",
|
|
side_effect=[
|
|
{"content": "第一页文字", "usage": {"total_tokens": 10}},
|
|
{"content": "第二页文字", "usage": {"total_tokens": 12}},
|
|
],
|
|
),
|
|
patch("services.pdf_ocr_service.reserve_avatar_tokens", return_value=reservation) as reserve,
|
|
patch("services.pdf_ocr_service.settle_reservation") as settle,
|
|
):
|
|
text = extract_scanned_pdf_text(
|
|
MagicMock(),
|
|
SimpleNamespace(id="avatar-1"),
|
|
"/tmp/scanned.pdf",
|
|
on_progress=lambda done, total: progress.append((done, total)),
|
|
)
|
|
|
|
assert text == "[第 1 页]\n第一页文字\n\n[第 2 页]\n第二页文字"
|
|
assert progress == [(1, 2), (2, 2)]
|
|
assert reserve.call_count == 2
|
|
assert settle.call_count == 2
|
|
|
|
|
|
def test_scanned_pdf_ocr_releases_tokens_after_retries_fail(monkeypatch):
|
|
fake_document = FakeDocument()
|
|
fake_document.page_count = 1
|
|
fake_pymupdf = SimpleNamespace(
|
|
open=lambda _path: fake_document,
|
|
Matrix=lambda x, y: (x, y),
|
|
csRGB="rgb",
|
|
)
|
|
monkeypatch.setitem(__import__("sys").modules, "pymupdf", fake_pymupdf)
|
|
monkeypatch.setenv("KNOWLEDGE_PDF_OCR_ATTEMPTS", "2")
|
|
reservation = SimpleNamespace()
|
|
config = SimpleNamespace(
|
|
api_key="configured",
|
|
ocr_model="qwen-vl-ocr",
|
|
vision_model="vision",
|
|
vision_max_tokens=2048,
|
|
)
|
|
|
|
with (
|
|
patch("services.pdf_ocr_service.get_chat_model_config", return_value=config),
|
|
patch("services.pdf_ocr_service.prepare_image", return_value=SimpleNamespace()),
|
|
patch("services.pdf_ocr_service.call_vision_model", side_effect=RuntimeError("timeout")) as call,
|
|
patch("services.pdf_ocr_service.reserve_avatar_tokens", return_value=reservation),
|
|
patch("services.pdf_ocr_service.release_reservation") as release,
|
|
patch("services.pdf_ocr_service.time.sleep"),
|
|
):
|
|
with pytest.raises(RuntimeError, match="第 1/1 页识别失败"):
|
|
extract_scanned_pdf_text(MagicMock(), SimpleNamespace(id="avatar-1"), "/tmp/scanned.pdf")
|
|
|
|
assert call.call_count == 2
|
|
release.assert_called_once()
|