DataOps-platform/tests/data_research/test_ocr_service.py

from __future__ import annotations

import pytest


class FakeProvider:
    def __init__(self, responses):
        self.responses = list(responses)
        self.calls = []

    def extract(self, image_bytes):
        self.calls.append(image_bytes)
        return self.responses.pop(0)


def block(text, bbox, confidence=0.9):
    from app.core.data_research.ocr.base import OcrBlock

    return OcrBlock(text=text, bbox=bbox, page=99, confidence=confidence)


def test_ocr_service_preserves_scanned_pdf_page_order_and_review_flags():
    from app.core.data_research.ocr.service import OcrService

    provider = FakeProvider(
        [
            [block("第一页", (0.1, 0.2, 0.8, 0.3), 0.95)],
            [block("第二页低置信", (0.0, 0.0, 1.0, 1.0), 0.4)],
        ]
    )
    results = OcrService(provider, review_threshold=0.75).extract_pages(
        [b"page-one", b"page-two"]
    )

    assert [result.page for result in results] == [1, 2]
    assert results[0].bbox == (0.1, 0.2, 0.8, 0.3)
    assert results[0].review_required is False
    assert results[1].review_required is True


def test_ocr_service_rejects_malformed_provider_response():
    from app.core.data_research.ocr.service import OcrResponseInvalid, OcrService

    provider = FakeProvider([[block("bad", (-1.0, 0.0, 2.0, 1.0))]])
    with pytest.raises(OcrResponseInvalid, match="bounding box"):
        OcrService(provider).extract_pages([b"image"])


def test_ocr_is_fail_closed_when_provider_is_disabled():
    from app.core.data_research.ocr.service import OcrDisabled, OcrService

    with pytest.raises(OcrDisabled, match="not configured"):
        OcrService(None).extract_pages([b"image"])


class FakeResponse:
    def __init__(self, payload, *, content=b"{}", headers=None):
        self._payload = payload
        self.content = content
        self.headers = headers or {}

    def raise_for_status(self):
        return None

    def json(self):
        return self._payload


def test_http_provider_enforces_timeout_tls_and_response_limit():
    from app.core.data_research.ocr.http_provider import HttpOcrProvider
    from app.core.data_research.ocr.service import OcrResponseInvalid

    calls = []

    def post(*args, **kwargs):
        calls.append((args, kwargs))
        return FakeResponse(
            {"blocks": [{"text": "编号", "bbox": [0, 0, 1, 1], "confidence": 0.9}]},
            content=b"x" * 20,
        )

    provider = HttpOcrProvider(
        "https://ocr.internal/v1/extract",
        token="secret-token",
        timeout_seconds=3,
        verify_tls=True,
        max_response_bytes=100,
        post=post,
    )
    blocks = provider.extract(b"png")
    assert blocks[0].text == "编号"
    assert calls[0][1]["timeout"] == 3
    assert calls[0][1]["verify"] is True

    provider.max_response_bytes = 10
    with pytest.raises(OcrResponseInvalid, match="response size") as error:
        provider.extract(b"png-secret-source")
    assert "secret-token" not in str(error.value)
    assert "png-secret-source" not in str(error.value)