xerrors--yuxi/backend/test/unit/knowledge/test_paddleocr_api_parser.py

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import pytest

import yuxi.knowledge.parser.paddleocr_api as paddleocr_api
from yuxi.knowledge.parser.base import DocumentParserException
from yuxi.knowledge.parser.paddleocr_api import PaddleOCRPPOCRv6Parser, PaddleOCRVLParser


@dataclass
class FakeResponse:
    status_code: int
    json_body: dict[str, Any] | None = None
    text: str = ""
    content: bytes = b""
    headers: dict[str, str] | None = None

    def json(self) -> dict[str, Any]:
        assert self.json_body is not None
        return self.json_body


def _build_file(tmp_path: Path, suffix: str = ".png") -> Path:
    file_path = tmp_path / f"sample{suffix}"
    file_path.write_bytes(b"fake image")
    return file_path


def test_paddleocr_vl_submits_model_specific_payload(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    file_path = _build_file(tmp_path)
    submitted: dict[str, Any] = {}

    def fake_post(url, headers=None, data=None, files=None, json=None, timeout=None):  # noqa: A002
        submitted["url"] = url
        submitted["data"] = data
        submitted["json"] = json
        submitted["files"] = files
        return FakeResponse(200, {"code": 0, "data": {"jobId": "job-vl"}})

    def fake_get(url, headers=None, timeout=None):
        if url.endswith("/job-vl"):
            return FakeResponse(200, {"data": {"state": "done", "resultUrl": {"jsonUrl": "https://result.test/vl"}}})
        row = {
            "result": {
                "layoutParsingResults": [
                    {"markdown": {"text": "VL markdown", "images": {}}, "outputImages": {"debug": "ignored"}}
                ]
            }
        }
        return FakeResponse(200, text=json.dumps(row))

    monkeypatch.setattr(paddleocr_api.requests, "post", fake_post)
    monkeypatch.setattr(paddleocr_api.requests, "get", fake_get)

    parser = PaddleOCRVLParser(api_token="token")
    result = parser.process_file(
        str(file_path),
        params={
            "optional_payload": {
                "useChartRecognition": True,
                "useTextlineOrientation": True,
                "ignored": True,
            }
        },
    )

    assert result == "VL markdown"
    assert submitted["data"]["model"] == "PaddleOCR-VL-1.6"
    optional_payload = json.loads(submitted["data"]["optionalPayload"])
    assert optional_payload == {
        "useDocOrientationClassify": False,
        "useDocUnwarping": False,
        "useChartRecognition": True,
    }


def test_paddleocr_pp_ocrv6_submits_model_specific_payload(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    file_path = _build_file(tmp_path)
    submitted: dict[str, Any] = {}

    def fake_post(url, headers=None, data=None, files=None, json=None, timeout=None):  # noqa: A002
        submitted["data"] = data
        return FakeResponse(200, {"data": {"jobId": "job-ocr"}})

    def fake_get(url, headers=None, timeout=None):
        if url.endswith("/job-ocr"):
            return FakeResponse(200, {"data": {"state": "done", "resultUrl": {"jsonUrl": "https://result.test/ocr"}}})
        row = {
            "result": {
                "ocrResults": [
                    {"prunedResult": {"rec_texts": ["PaddleOCR API Test", "", "Invoice total: 123.45"]}}
                ]
            }
        }
        return FakeResponse(200, text=json.dumps(row))

    monkeypatch.setattr(paddleocr_api.requests, "post", fake_post)
    monkeypatch.setattr(paddleocr_api.requests, "get", fake_get)

    parser = PaddleOCRPPOCRv6Parser(api_token="token")
    result = parser.process_file(
        str(file_path),
        params={
            "optional_payload": {
                "useTextlineOrientation": True,
                "useChartRecognition": True,
                "ignored": True,
            }
        },
    )

    assert result == "PaddleOCR API Test\nInvoice total: 123.45"
    assert submitted["data"]["model"] == "PP-OCRv6"
    optional_payload = json.loads(submitted["data"]["optionalPayload"])
    assert optional_payload == {
        "useDocOrientationClassify": False,
        "useDocUnwarping": False,
        "useTextlineOrientation": True,
    }


def test_paddleocr_url_input_uses_json_payload(monkeypatch: pytest.MonkeyPatch) -> None:
    submitted: dict[str, Any] = {}

    def fake_post(url, headers=None, data=None, files=None, json=None, timeout=None):  # noqa: A002
        submitted["headers"] = headers
        submitted["data"] = data
        submitted["files"] = files
        submitted["json"] = json
        return FakeResponse(200, {"data": {"jobId": "job-url"}})

    def fake_get(url, headers=None, timeout=None):
        if url.endswith("/job-url"):
            return FakeResponse(200, {"data": {"state": "done", "resultUrl": {"jsonUrl": "https://result.test/url"}}})
        row = {"result": {"ocrResults": [{"prunedResult": {"rec_texts": ["url text"]}}]}}
        return FakeResponse(200, text=json.dumps(row))

    monkeypatch.setattr(paddleocr_api.requests, "post", fake_post)
    monkeypatch.setattr(paddleocr_api.requests, "get", fake_get)

    parser = PaddleOCRPPOCRv6Parser(api_token="token")
    result = parser.process_file("https://example.test/file.png?signature=abc")

    assert result == "url text"
    assert submitted["headers"]["Content-Type"] == "application/json"
    assert submitted["data"] is None
    assert submitted["files"] is None
    assert submitted["json"]["fileUrl"] == "https://example.test/file.png?signature=abc"
    assert submitted["json"]["model"] == "PP-OCRv6"


def test_paddleocr_poll_handles_pending_running_and_done(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    file_path = _build_file(tmp_path)
    states = iter(["pending", "running", "done"])
    sleep_calls: list[float] = []

    monkeypatch.setattr(
        paddleocr_api.requests,
        "post",
        lambda *args, **kwargs: FakeResponse(200, {"data": {"jobId": "job-poll"}}),
    )

    def fake_get(url, headers=None, timeout=None):
        if url.endswith("/job-poll"):
            state = next(states)
            data: dict[str, Any] = {"state": state}
            if state == "done":
                data["resultUrl"] = {"jsonUrl": "https://result.test/poll"}
            return FakeResponse(200, {"data": data})
        row = {"result": {"ocrResults": [{"prunedResult": {"rec_texts": ["done text"]}}]}}
        return FakeResponse(200, text=json.dumps(row))

    monkeypatch.setattr(paddleocr_api.requests, "get", fake_get)
    monkeypatch.setattr(paddleocr_api.time, "sleep", lambda seconds: sleep_calls.append(seconds))

    parser = PaddleOCRPPOCRv6Parser(api_token="token")
    result = parser.process_file(str(file_path), params={"poll_interval_seconds": 0.25})

    assert result == "done text"
    assert sleep_calls == [0.25, 0.25]


def test_paddleocr_failed_job_raises_parser_exception(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    file_path = _build_file(tmp_path)

    monkeypatch.setattr(
        paddleocr_api.requests,
        "post",
        lambda *args, **kwargs: FakeResponse(200, {"data": {"jobId": "job-failed"}}),
    )
    monkeypatch.setattr(
        paddleocr_api.requests,
        "get",
        lambda *args, **kwargs: FakeResponse(
            200,
            {"data": {"state": "failed", "errorMsg": "quota exceeded"}},
        ),
    )

    parser = PaddleOCRVLParser(api_token="token")
    with pytest.raises(DocumentParserException, match="quota exceeded"):
        parser.process_file(str(file_path))


def test_paddleocr_missing_token_health_and_parse_error(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    file_path = _build_file(tmp_path)
    monkeypatch.delenv("PADDLEOCR_API_TOKEN", raising=False)
    parser = PaddleOCRVLParser()

    health = parser.check_health()

    assert health["status"] == "unavailable"
    assert "PADDLEOCR_API_TOKEN" in health["message"]
    with pytest.raises(DocumentParserException, match="PADDLEOCR_API_TOKEN"):
        parser.process_file(str(file_path))


def test_paddleocr_configured_token_health_does_not_submit_job() -> None:
    parser = PaddleOCRVLParser(api_token="token")

    health = parser.check_health()

    assert health["status"] == "configured"
    assert "解析时验证" in health["message"]
    assert health["details"]["model"] == "PaddleOCR-VL-1.6"


def test_paddleocr_vl_uploads_markdown_images(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    file_path = _build_file(tmp_path)
    uploaded: dict[str, Any] = {}

    monkeypatch.setattr(
        paddleocr_api.requests,
        "post",
        lambda *args, **kwargs: FakeResponse(200, {"data": {"jobId": "job-images"}}),
    )

    def fake_get(url, headers=None, timeout=None):
        if url.endswith("/job-images"):
            return FakeResponse(
                200,
                {"data": {"state": "done", "resultUrl": {"jsonUrl": "https://result.test/images"}}},
            )
        if url == "https://image.test/table.png":
            return FakeResponse(200, content=b"image-bytes", headers={"Content-Type": "image/png"})
        row = {
            "result": {
                "layoutParsingResults": [
                    {
                        "markdown": {
                            "text": "before ![](images/table.png) after",
                            "images": {"images/table.png": "https://image.test/table.png"},
                        }
                    }
                ]
            }
        }
        return FakeResponse(200, text=json.dumps(row))

    class FakeMinioClient:
        def ensure_bucket_exists(self, bucket_name):
            uploaded["bucket_name"] = bucket_name

        def upload_file(self, bucket_name, object_name, data):
            uploaded["object_name"] = object_name
            uploaded["data"] = data
            return type("UploadResult", (), {"url": "minio://public/kb/table.png"})()

    monkeypatch.setattr(paddleocr_api.requests, "get", fake_get)
    monkeypatch.setattr(paddleocr_api, "get_minio_client", lambda: FakeMinioClient())
    monkeypatch.setattr(paddleocr_api.time, "time", lambda: 1.0)

    parser = PaddleOCRVLParser(api_token="token")
    result = parser.process_file(
        str(file_path),
        params={"image_bucket": "public", "image_prefix": "kb/images"},
    )

    assert result == "before ![](minio://public/kb/table.png) after"
    assert uploaded == {
        "bucket_name": "public",
        "object_name": "kb/images/1000000_table.png",
        "data": b"image-bytes",
    }