83 lines
2.6 KiB
Python
83 lines
2.6 KiB
Python
"""Tests for OCR-augmented extraction (rex/comments — #663)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import fitz
|
|
|
|
from rex.comments.combine import extract_comment, parse_combined
|
|
from rex.comments.extract import extract_attachment
|
|
|
|
|
|
def _fake_engine(text: str):
|
|
def engine(path: Path) -> str: # noqa: ARG001
|
|
return text
|
|
|
|
return engine
|
|
|
|
|
|
def _write_image(path: Path) -> None:
|
|
# 1x1 PNG via pymupdf — content is irrelevant, the fake engine answers.
|
|
pix = fitz.Pixmap(fitz.csRGB, fitz.IRect(0, 0, 8, 8))
|
|
pix.save(str(path))
|
|
|
|
|
|
def _write_scanned_pdf(path: Path) -> None:
|
|
# a PDF with no text layer
|
|
doc = fitz.open()
|
|
doc.new_page()
|
|
doc.save(str(path))
|
|
doc.close()
|
|
|
|
|
|
class TestExtractAttachmentOcr:
|
|
def test_image_without_engine_stays_queued(self, tmp_path):
|
|
img = tmp_path / "a.png"
|
|
_write_image(img)
|
|
r = extract_attachment(img)
|
|
assert r.status == "ocr_needed"
|
|
|
|
def test_image_with_engine_extracts(self, tmp_path):
|
|
img = tmp_path / "a.png"
|
|
_write_image(img)
|
|
r = extract_attachment(img, ocr_engine=_fake_engine("Dear CMS, " * 20))
|
|
assert r.status == "ok"
|
|
assert "Dear CMS" in r.text
|
|
|
|
def test_scanned_pdf_with_engine_extracts(self, tmp_path):
|
|
pdf = tmp_path / "scan.pdf"
|
|
_write_scanned_pdf(pdf)
|
|
r = extract_attachment(
|
|
pdf, ocr_engine=_fake_engine("Scanned letter text " * 10)
|
|
)
|
|
assert r.status == "ok"
|
|
assert "Scanned letter" in r.text
|
|
|
|
def test_engine_short_output_stays_queued(self, tmp_path):
|
|
img = tmp_path / "a.png"
|
|
_write_image(img)
|
|
r = extract_attachment(img, ocr_engine=_fake_engine("x"))
|
|
assert r.status == "ocr_needed"
|
|
|
|
|
|
class TestExtractCommentOcrRerun:
|
|
def test_force_rerun_with_engine_updates_combined(self, tmp_path):
|
|
d = tmp_path / "CMS-2023-0121" / "CMS-2023-0121-9999"
|
|
d.mkdir(parents=True)
|
|
_write_image(d / "attachment_1.png")
|
|
# first pass without OCR → ocr_needed in frontmatter
|
|
out = extract_comment(d, inline_body="inline")
|
|
fm, _ = parse_combined(out.read_text())
|
|
assert fm["attachments"][0]["status"] == "ocr_needed"
|
|
# OCR pass
|
|
out2 = extract_comment(
|
|
d,
|
|
inline_body="inline",
|
|
force=True,
|
|
ocr_engine=_fake_engine("Recovered scanned palliative text " * 5),
|
|
)
|
|
fm2, body2 = parse_combined(out2.read_text())
|
|
assert fm2["attachments"][0]["status"] == "ok"
|
|
assert "Recovered scanned palliative text" in body2
|