office_translator/tests/test_scanned_pdf_ocr.py

"""Scanned PDF detection and the Mistral OCR translation path."""

import fitz
import pytest

from config import config
from services.mistral_ocr import MistralOCRClient
from services.providers.base import TranslationProvider
from services.providers.schemas import TranslationRequest, TranslationResponse
from translators.pdf_translator import PDFTranslator


FAKE_OCR_PAGES = [
    "# Facture\n\nMontant total : 1 250 euros\n\n![logo](img.png)",
    "Le paiement est du sous 30 jours.",
]

# Minimal French → English mapping for the fake provider; anything else is
# prefixed so tests can assert replacement happened.
TRANSLATIONS = {
    "Facture": "Invoice",
    "Montant total : 1 250 euros": "Total amount: 1,250 euros",
    "Le paiement est du sous 30 jours.": "Payment is due within 30 days.",
}


class FakeProvider(TranslationProvider):
    """New-style provider with canned translations (no network)."""

    def get_name(self) -> str:
        return "fake"

    def is_available(self) -> bool:
        return True

    def translate_text(self, request: TranslationRequest) -> TranslationResponse:
        # Pages arrive as multi-line text: translate line by line so the
        # canned mapping applies regardless of how lines are grouped.
        lines = []
        for line in request.text.split("\n"):
            stripped = line.strip()
            if not stripped:
                lines.append(line)
                continue
            lines.append(TRANSLATIONS.get(stripped, f"[EN] {stripped}"))
        return TranslationResponse(
            translated_text="\n".join(lines),
            provider_name=self.get_name(),
            from_cache=False,
        )


def _make_scanned_pdf(path):
    """Image-only PDF: one page almost fully covered by a picture, no text."""
    pix = fitz.Pixmap(fitz.csRGB, fitz.IRect(0, 0, 10, 10))
    pix.clear_with(90)
    png_bytes = pix.tobytes("png")
    doc = fitz.open()
    page = doc.new_page(width=612, height=792)
    page.insert_image(fitz.Rect(20, 20, 592, 772), stream=png_bytes)
    doc.save(str(path))
    doc.close()
    return path


def _make_text_pdf(path):
    doc = fitz.open()
    page = doc.new_page()
    page.insert_text(
        (72, 100),
        "Real selectable text with plenty of characters to exceed the scanned threshold. " * 2,
        fontsize=11,
    )
    doc.save(str(path))
    doc.close()
    return path


@pytest.fixture
def no_api_key(monkeypatch):
    monkeypatch.delenv("MISTRAL_API_KEY", raising=False)
    monkeypatch.setattr(config, "MISTRAL_API_KEY", "")
    monkeypatch.setattr(config, "MISTRAL_OCR_ENABLED", True)


@pytest.fixture
def api_key(monkeypatch):
    # Set BOTH the env var and the config attribute: some other test in the
    # full suite reloads the config module, and pdf_translator imports
    # `config` at call time — after a reload only the env var is still
    # visible (the class attributes are re-read from the environment).
    monkeypatch.setenv("MISTRAL_API_KEY", "test-key")
    monkeypatch.setattr(config, "MISTRAL_API_KEY", "test-key")
    monkeypatch.setattr(config, "MISTRAL_OCR_ENABLED", True)


class TestScannedDetection:
    def test_image_only_pdf_is_scanned(self, tmp_path):
        pdf = _make_scanned_pdf(tmp_path / "scan.pdf")
        translator = PDFTranslator(provider=None)
        assert translator._is_scanned_pdf(pdf) is True

    def test_text_pdf_is_not_scanned(self, tmp_path):
        pdf = _make_text_pdf(tmp_path / "text.pdf")
        translator = PDFTranslator(provider=None)
        assert translator._is_scanned_pdf(pdf) is False

    def test_title_only_pdf_is_not_scanned(self, tmp_path):
        """A sparse but textual page (no raster image) stays on the layout path."""
        doc = fitz.open()
        page = doc.new_page()
        page.insert_text((72, 100), "Facture 2026", fontsize=18)
        doc.save(str(tmp_path / "title.pdf"))
        doc.close()
        translator = PDFTranslator(provider=None)
        assert translator._is_scanned_pdf(tmp_path / "title.pdf") is False


class TestMarkdownCleanup:
    def test_images_links_headings_removed(self):
        md = "# Titre\n\n![img](x.png)\n\n[lien](http://x) texte"
        out = PDFTranslator._markdown_to_text(md)
        assert "Titre" in out
        assert "![img]" not in out
        assert "(http://x)" not in out
        assert "#" not in out
        assert "lien texte" in out

    def test_table_pipes_removed(self):
        out = PDFTranslator._markdown_to_text("| A | B |\n|---|---|\n| un | deux |")
        assert "|" not in out
        assert "un" in out and "deux" in out


class TestScannedOCRPath:
    def test_translate_scanned_pdf_end_to_end(self, tmp_path, api_key, monkeypatch):
        pdf = _make_scanned_pdf(tmp_path / "scan.pdf")
        monkeypatch.setattr(
            MistralOCRClient,
            "extract_pdf_text",
            lambda self, path, progress_callback=None: list(FAKE_OCR_PAGES),
        )

        translator = PDFTranslator(provider=FakeProvider())
        # Configure OCR the way the production route does (set_ocr_config).
        # Going through the config module is fragile in the full suite:
        # another test replaces sys.modules["config"], so a config-module
        # patch may never reach the translator.
        translator.set_ocr_config(api_key="test-key", enabled=True)
        out = tmp_path / "out.pdf"
        result = translator.translate_file(pdf, out, "en", "fr")

        assert result.exists() and result.stat().st_size > 0
        doc = fitz.open(str(result))
        text = "\n".join(p.get_text("text") for p in doc)
        doc.close()

        assert "Invoice" in text
        assert "Payment is due within 30 days" in text
        assert "Facture" not in text
        assert "logo" not in text          # markdown images stripped
        assert "|" not in text             # markdown tables stripped

    def test_missing_api_key_raises_clear_error(self, tmp_path, no_api_key):
        pdf = _make_scanned_pdf(tmp_path / "scan.pdf")
        translator = PDFTranslator(provider=FakeProvider())
        with pytest.raises(RuntimeError) as exc:
            translator.translate_file(pdf, tmp_path / "out.pdf", "en", "fr")
        assert "MISTRAL_API_KEY" in str(exc.value)

    def test_disabled_ocr_raises_clear_error(self, tmp_path, monkeypatch):
        monkeypatch.setattr(config, "MISTRAL_API_KEY", "test-key")
        monkeypatch.setattr(config, "MISTRAL_OCR_ENABLED", False)
        pdf = _make_scanned_pdf(tmp_path / "scan.pdf")
        translator = PDFTranslator(provider=FakeProvider())
        with pytest.raises(RuntimeError):
            translator.translate_file(pdf, tmp_path / "out.pdf", "en", "fr")

    def test_text_pdf_bypasses_ocr(self, tmp_path, api_key, monkeypatch):
        pdf = _make_text_pdf(tmp_path / "text.pdf")
        called = {"ocr": False}

        def _fail(self, path, progress_callback=None):
            called["ocr"] = True
            return []

        monkeypatch.setattr(MistralOCRClient, "extract_pdf_text", _fail)

        translator = PDFTranslator(provider=FakeProvider())
        out = tmp_path / "out.pdf"
        translator.translate_file(pdf, out, "en", "fr")
        assert called["ocr"] is False