DataOps-platform/tests/data_research/test_image_extractor.py

from __future__ import annotations

from app.core.data_research.extractors.base import ExtractionContext


class FakeOcrService:
    def extract_pages(self, pages):
        from app.core.data_research.ocr.service import OcrResult

        assert pages == [b"image"]
        return [
            OcrResult(
                text="customer_id 客户编号",
                bbox=(0.1, 0.2, 0.7, 0.3),
                page=1,
                confidence=0.65,
                review_required=True,
            )
        ]


def test_png_and_jpeg_are_supported_with_normalized_ocr_evidence():
    from app.core.data_research.extractors.image import ImageExtractor

    extractor = ImageExtractor(FakeOcrService())
    assert extractor.can_handle("image/png", "scan.png")
    assert extractor.can_handle("image/jpeg", "scan.jpg")

    batch = extractor.extract(
        b"image",
        ExtractionContext(filename="scan.png", media_type="image/png"),
    )
    assert batch.parser_version == "image-ocr-v1"
    assert batch.items[0].data["review_required"] is True
    assert batch.items[0].evidence[0].locator == {
        "kind": "image.ocr",
        "page": 1,
        "bbox": [0.1, 0.2, 0.7, 0.3],
    }
    assert batch.items[0].evidence[0].confidence == 0.65