rag-service/ocr-service/tests/test_render.py

import hashlib
import threading
import time
import sys
from pathlib import Path

import pytest

sys.path.insert(0, str(Path(__file__).parents[1]))

from app.engine import EngineLine, PaddleOcrEngine
from app import render
from app.render import PdfRenderError, process_pdf, render_pdf_pages


FIXTURE = Path(__file__).parents[2] / "tests" / "fixtures" / "ocr" / "native-three-pages.pdf"
DOCUMENT_SHA256 = hashlib.sha256(FIXTURE.read_bytes()).hexdigest()


class DeterministicEngine:
    def __init__(self) -> None:
        self.calls = 0

    def recognize(self, _image: object) -> list[EngineLine]:
        self.calls += 1
        if self.calls % 2 == 1:
            return [
                EngineLine("FATo7", 0.98, (120, 340, 245, 372)),
                EngineLine("second line", 0.74, (80, 410, 300, 450)),
            ]
        return [EngineLine("page three", 0.91, (50, 60, 250, 100))]


def test_render_produces_selected_200_dpi_png_pages() -> None:
    pages = render_pdf_pages(FIXTURE.read_bytes(), [1, 3])

    assert [(page.page, page.width, page.height) for page in pages] == [
        (1, 1700, 2200),
        (3, 1700, 2200),
    ]
    assert all(page.png.startswith(b"\x89PNG\r\n\x1a\n") for page in pages)


def test_render_rejects_invalid_pages_and_pixel_limit_before_image_creation() -> None:
    with pytest.raises(PdfRenderError, match="one-based"):
        render_pdf_pages(FIXTURE.read_bytes(), [0])
    with pytest.raises(PdfRenderError, match="25 megapixels"):
        render_pdf_pages(FIXTURE.read_bytes(), [1], max_pixels=1_000_000)


def test_render_serializes_pdfium_calls_with_the_process_lock(monkeypatch: pytest.MonkeyPatch) -> None:
    active = 0
    peak = 0

    class Document:
        def __init__(self, _pdf: bytes) -> None:
            nonlocal active, peak
            active += 1
            peak = max(peak, active)
            time.sleep(0.03)
            active -= 1
            raise RuntimeError("test document")

    class Pdfium:
        PdfDocument = Document

    monkeypatch.setitem(sys.modules, "pypdfium2", Pdfium())
    def invoke() -> None:
        with pytest.raises(PdfRenderError):
            render_pdf_pages(FIXTURE.read_bytes(), [1])

    threads = [threading.Thread(target=invoke) for _ in range(2)]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join()
    assert peak == 1


def test_render_builds_repeatable_result_schema_with_deterministic_engine() -> None:
    def run() -> dict:
        return process_pdf(
            job_id="ocr_repeatable",
            document_sha256=DOCUMENT_SHA256,
            pdf=FIXTURE.read_bytes(),
            requested_pages=[1, 3],
            engine=DeterministicEngine(),
            processing_ms=lambda _page: 17,
        )

    first = run()
    assert first == run()
    assert first["schemaVersion"] == "1"
    assert first["documentSha256"] == DOCUMENT_SHA256
    assert 0 < first["pages"][0]["metrics"]["inkCoverage"] < 1
    assert first["engine"] == {
        "name": "paddleocr",
        "version": "3.4.0",
        "runtime": "paddlepaddle-3.2.2",
        "device": "cpu",
        "configVersion": "ocr-v1",
        "dpi": 200,
    }
    assert first["pages"][0] == {
        "page": 1,
        "width": 1700,
        "height": 2200,
        "processingMs": 17,
        "text": "FATo7\nsecond line",
        "metrics": {
            "lineCount": 2,
            "nonWhitespaceCharacters": 15,
            "inkCoverage": first["pages"][0]["metrics"]["inkCoverage"],
            "medianConfidence": 0.86,
            "p10Confidence": 0.74,
            "lowConfidenceLineRatio": 0.0,
        },
        "lines": [
            {
                "lineId": "p1-l01-120-340-245-372",
                "text": "FATo7",
                "confidence": 0.98,
                "bbox": [120, 340, 245, 372],
            },
            {
                "lineId": "p1-l02-80-410-300-450",
                "text": "second line",
                "confidence": 0.74,
                "bbox": [80, 410, 300, 450],
            },
        ],
    }
    assert [page["page"] for page in first["pages"]] == [1, 3]


def test_render_paddle_adapter_preserves_text_confidence_and_boxes() -> None:
    class Prediction:
        json = {
            "res": {
                "rec_texts": ["CBGO4a", "NSAvo6"],
                "rec_scores": [0.97, 0.83],
                "rec_boxes": [[10, 20, 110, 50], [15, 70, 130, 100]],
            }
        }

    class Pipeline:
        def predict(self, _image: object) -> list[Prediction]:
            return [Prediction()]

    lines = PaddleOcrEngine(pipeline=Pipeline()).recognize(object())

    assert lines == [
        EngineLine("CBGO4a", 0.97, (10, 20, 110, 50)),
        EngineLine("NSAvo6", 0.83, (15, 70, 130, 100)),
    ]


def test_render_container_pins_cpu_runtime_models_and_single_worker() -> None:
    service_root = Path(__file__).parents[1]
    requirements = (service_root / "requirements.txt").read_text()
    dockerfile = (service_root / "Dockerfile").read_text()
    dockerignore = (service_root / "Dockerfile.dockerignore").read_text().splitlines()
    model_loader = (service_root / "app" / "models.py").read_text()

    assert "paddleocr==3.4.0" in requirements
    assert "paddlepaddle==3.2.2" in requirements
    assert "pypdfium2==4.30.0" in requirements
    assert "RUN python -m app.models" in dockerfile
    assert 'OCR_JOBS_DB="/data/jobs/jobs.db"' in dockerfile
    assert '"--workers", "1"' in dockerfile
    assert "USER ocr" in dockerfile
    assert 'resource.cpu.max="3"' in dockerfile
    assert 'resource.memory.max="5GiB"' in dockerfile
    assert "PaddleOcrEngine" in model_loader
    assert dockerignore == [
        "**",
        "!ocr-service/",
        "!ocr-service/app/",
        "!ocr-service/app/**",
        "!ocr-service/requirements.txt",
    ]