116 lines
4.4 KiB
Python
116 lines
4.4 KiB
Python
"""Unit tests for ExtractorFactory + detect_format + the OCR-fallback flow."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from contract_check.core.extraction import (
|
|
UnsupportedFormatError,
|
|
detect_format,
|
|
extract_document,
|
|
get_factory,
|
|
)
|
|
from tests.unit.test_extraction_adapters import PARAGRAPH, build_docx_bytes, build_pdf_bytes
|
|
|
|
# ── detect_format: suffix / mime / magic precedence ──────────────────────────
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("filename", "expected"),
|
|
[
|
|
("contract.pdf", "pdf"),
|
|
("contract.PDF", "pdf"),
|
|
("contract.docx", "docx"),
|
|
("contract.rtf", "rtf"),
|
|
("contract.txt", "txt"),
|
|
("data.csv", "txt"),
|
|
("scan.png", "image"),
|
|
("scan.jpg", "image"),
|
|
("scan.jpeg", "image"),
|
|
("scan.tiff", "image"),
|
|
],
|
|
)
|
|
def test_detect_by_suffix(filename: str, expected: str) -> None:
|
|
assert detect_format(b"", filename=filename) == expected
|
|
|
|
|
|
def test_detect_by_mime_when_suffix_unknown() -> None:
|
|
assert detect_format(b"", mime="application/pdf", filename="contract") == "pdf"
|
|
assert detect_format(b"", mime="image/png", filename="photo") == "image"
|
|
|
|
|
|
def test_magic_beats_wrong_suffix() -> None:
|
|
# A real PDF mislabeled .docx must route to the PDF adapter.
|
|
pytest.importorskip("magic")
|
|
assert detect_format(build_pdf_bytes(), mime="", filename="contract.docx") == "pdf"
|
|
|
|
|
|
def test_unsupported_format_raises() -> None:
|
|
with pytest.raises(UnsupportedFormatError):
|
|
detect_format(b"", filename="archive.zip")
|
|
binary_soup = bytes(range(256)) * 4 # deliberately not decodable text
|
|
with pytest.raises(UnsupportedFormatError):
|
|
detect_format(binary_soup, mime="application/x-msdownload", filename="doc.exe")
|
|
|
|
|
|
def test_docx_zip_magic_does_not_shadow_suffix() -> None:
|
|
# DOCX is a zip; python-magic reports application/zip which the factory
|
|
# deliberately does not map — the suffix must decide.
|
|
pytest.importorskip("magic")
|
|
assert detect_format(build_docx_bytes(), filename="contract.docx") == "docx"
|
|
|
|
|
|
# ── factory ──────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_factory_returns_per_format_adapters() -> None:
|
|
factory = get_factory()
|
|
assert type(factory.get_extractor(b"", filename="a.pdf")).__name__ == "PyMuPDFExtractor"
|
|
assert type(factory.get_extractor(b"", filename="a.docx")).__name__ == "MammothDocxExtractor"
|
|
assert type(factory.get_extractor(b"", filename="a.rtf")).__name__ == "RtfExtractor"
|
|
assert type(factory.get_extractor(b"", filename="a.txt")).__name__ == "TxtExtractor"
|
|
assert type(factory.get_extractor(b"", filename="a.png")).__name__ == "TesseractOcrExtractor"
|
|
|
|
|
|
# ── extract_document: end-to-end over bytes ──────────────────────────────────
|
|
|
|
|
|
def test_extract_document_pdf() -> None:
|
|
result = extract_document(build_pdf_bytes(), filename="c.pdf")
|
|
assert result.metadata["format"] == "pdf"
|
|
assert "арбитражном суде" in result.markdown
|
|
|
|
|
|
def test_extract_document_docx() -> None:
|
|
result = extract_document(build_docx_bytes(), filename="c.docx")
|
|
assert result.markdown.startswith("# ")
|
|
assert result.is_structured is True
|
|
|
|
|
|
def test_extract_document_txt_cp1251() -> None:
|
|
result = extract_document((PARAGRAPH * 10).encode("windows-1251"), filename="c.txt")
|
|
assert "Стороны" in result.markdown
|
|
assert result.metadata["format"] == "txt"
|
|
|
|
|
|
def test_extract_document_unsupported() -> None:
|
|
with pytest.raises(UnsupportedFormatError):
|
|
extract_document(b"PK\x03\x04 whatever", filename="c.zip")
|
|
|
|
|
|
def test_extract_document_pdf_scan_falls_back_to_ocr() -> None:
|
|
# A PDF with no text layer triggers the OCR fallback inside the
|
|
# orchestrator. Without the tesseract engine the fallback raises
|
|
# OCRError (retryable) — with the engine, the integration suite covers
|
|
# the happy path.
|
|
import pymupdf
|
|
|
|
from contract_check.core.analysis.ocr import OCRError
|
|
from contract_check.core.extraction import ExtractionFailedError
|
|
|
|
doc = pymupdf.open()
|
|
doc.new_page() # blank page: no text at all
|
|
data = doc.tobytes()
|
|
doc.close()
|
|
with pytest.raises((OCRError, ExtractionFailedError)):
|
|
extract_document(data, filename="scan.pdf")
|