doc_processer/tests/api/v1/endpoints/test_image_endpoint.py

import numpy as np
from fastapi import FastAPI
from fastapi.testclient import TestClient

from app.api.v1.endpoints.image import router
from app.core.dependencies import get_glmocr_endtoend_service, get_image_processor


class _FakeImageProcessor:
    def preprocess(self, image_url=None, image_base64=None):
        return np.zeros((8, 8, 3), dtype=np.uint8)


class _FakeOCRService:
    def __init__(self, result=None, error=None):
        self._result = result or {"markdown": "md", "latex": "tex", "mathml": "mml", "mml": "xml"}
        self._error = error

    def recognize(self, image):
        if self._error:
            raise self._error
        return self._result


def _build_client(image_processor=None, ocr_service=None):
    app = FastAPI()
    app.include_router(router)
    app.dependency_overrides[get_image_processor] = lambda: image_processor or _FakeImageProcessor()
    app.dependency_overrides[get_glmocr_endtoend_service] = lambda: ocr_service or _FakeOCRService()
    return TestClient(app)


def test_image_endpoint_requires_exactly_one_of_image_url_or_image_base64():
    client = _build_client()

    missing = client.post("/ocr", json={})
    both = client.post(
        "/ocr", json={"image_url": "https://example.com/a.png", "image_base64": "abc"}
    )

    assert missing.status_code == 422
    assert both.status_code == 422


def test_image_endpoint_returns_503_for_runtime_error():
    client = _build_client(ocr_service=_FakeOCRService(error=RuntimeError("backend unavailable")))

    response = client.post("/ocr", json={"image_url": "https://example.com/a.png"})

    assert response.status_code == 503
    assert response.json()["detail"] == "backend unavailable"


def test_image_endpoint_returns_500_for_unexpected_error():
    client = _build_client(ocr_service=_FakeOCRService(error=ValueError("boom")))

    response = client.post("/ocr", json={"image_url": "https://example.com/a.png"})

    assert response.status_code == 500
    assert response.json()["detail"] == "Internal server error"


def test_image_endpoint_returns_ocr_payload():
    client = _build_client()

    response = client.post("/ocr", json={"image_base64": "ZmFrZQ=="})

    assert response.status_code == 200
    assert response.json() == {
        "latex": "tex",
        "markdown": "md",
        "mathml": "mml",
        "mml": "xml",
        "layout_info": {"regions": [], "MixedRecognition": False},
        "recognition_mode": "",
    }


def test_image_endpoint_real_e2e_with_env_services():
    from app.main import app

    image_url = (
        "https://static.texpixel.com/formula/012dab3e-fb31-4ecd-90fc-6957458ee309.png"
        "?Expires=1773049821&OSSAccessKeyId=TMP.3KnrJUz7aXHoU9rLTAih4MAyPGd9zyGRHiqg9AyH6TY6NKtzqT2yr4qo7Vwf8fMRFCBrWXiCFrbBwC3vn7U6mspV2NeU1K"
        "&Signature=oynhP0OLIgFI0Sv3z2CWeHPT2Ck%3D"
    )

    with TestClient(app) as client:
        response = client.post(
            "/doc_process/v1/image/ocr",
            json={"image_url": image_url},
            headers={"x-request-id": "test-e2e"},
        )

    assert response.status_code == 200, response.text
    payload = response.json()
    assert isinstance(payload["markdown"], str)
    assert payload["markdown"].strip()
    assert set(payload) >= {"markdown", "latex", "mathml", "mml"}