tests/test_parser_client.py

from __future__ import annotations

import base64

import pytest
import requests

from parser_client import (
    ParserClient,
    ParserConfig,
    ParserServiceError,
    api_key_name_for_environment,
    build_payload,
    selected_environment,
)


class FakeResponse:
    def __init__(
        self,
        status: int,
        body: dict,
        headers: dict | None = None,
    ) -> None:
        self.status_code = status
        self._body = body
        self.headers = headers or {}
        self.ok = 200 <= status < 300

    def json(self) -> dict:
        return self._body


class FakeSession:
    def __init__(self, responses: list[FakeResponse | Exception]) -> None:
        self.responses = list(responses)
        self.calls: list[dict] = []

    def post(self, url: str, **kwargs):
        self.calls.append({"url": url, **kwargs})
        response = self.responses.pop(0)
        if isinstance(response, Exception):
            raise response
        return response


def config() -> ParserConfig:
    return ParserConfig(
        api_url="https://stg.api.cohere.test/v2/parse",
        api_key="secret-key",
        model="parse-v5.0",
        timeout=12,
        max_retries=3,
    )


def clear_cohere_environment(monkeypatch: pytest.MonkeyPatch) -> None:
    for name in (
        "COHERE_ENV",
        "COHERE_API_KEY",
        "COHERE_STG_API_KEY",
        "COHERE_API_URL",
    ):
        monkeypatch.delenv(name, raising=False)


def test_environment_defaults_to_staging(monkeypatch: pytest.MonkeyPatch) -> None:
    clear_cohere_environment(monkeypatch)
    monkeypatch.setenv("COHERE_STG_API_KEY", "staging-key")

    parsed = ParserConfig.from_environment()

    assert selected_environment() == "staging"
    assert api_key_name_for_environment() == "COHERE_STG_API_KEY"
    assert parsed.api_key == "staging-key"
    assert parsed.api_url == "https://stg.api.cohere.ai/v2/parse"


def test_production_environment_uses_production_key_and_url(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    clear_cohere_environment(monkeypatch)
    monkeypatch.setenv("COHERE_ENV", "production")
    monkeypatch.setenv("COHERE_API_KEY", "production-key")
    monkeypatch.setenv("COHERE_STG_API_KEY", "staging-key")

    parsed = ParserConfig.from_environment()

    assert api_key_name_for_environment() == "COHERE_API_KEY"
    assert parsed.api_key == "production-key"
    assert parsed.api_url == "https://api.cohere.com/v2/parse"


def test_api_url_can_override_the_selected_environment(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    clear_cohere_environment(monkeypatch)
    monkeypatch.setenv("COHERE_STG_API_KEY", "staging-key")
    monkeypatch.setenv("COHERE_API_URL", "https://gateway.example.test/v2/parse")

    parsed = ParserConfig.from_environment()

    assert parsed.api_url == "https://gateway.example.test/v2/parse"


def test_environment_rejects_invalid_values(monkeypatch: pytest.MonkeyPatch) -> None:
    clear_cohere_environment(monkeypatch)
    monkeypatch.setenv("COHERE_ENV", "preview")

    with pytest.raises(ParserServiceError, match="staging.*production"):
        ParserConfig.from_environment()


def test_environment_requires_its_matching_key(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    clear_cohere_environment(monkeypatch)
    monkeypatch.setenv("COHERE_ENV", "production")
    monkeypatch.setenv("COHERE_STG_API_KEY", "staging-key")

    with pytest.raises(ParserServiceError, match="configured for production"):
        ParserConfig.from_environment()


def test_build_payload_uses_image_only_contract() -> None:
    payload = build_payload("parser-model", b"png-bytes")

    assert payload == {
        "model": "parser-model",
        "document": {
            "type": "image_url",
            "image_url": payload["document"]["image_url"],
        },
        "output_format": "blocks",
    }
    data_url = payload["document"]["image_url"]
    assert data_url.startswith("data:image/png;base64,")
    assert base64.b64decode(data_url.split(",", 1)[1]) == b"png-bytes"


def parse_response(content: str = "# Invoice") -> dict:
    return {
        "id": "parse-response-123",
        "pages": [
            {
                "index": 0,
                "type": "blocks",
                "blocks": [
                    {"type": "text", "text": {"content": content}},
                    {
                        "type": "table",
                        "table": {
                            "type": "html",
                            "html": "<table><tr><td>Total</td></tr></table>",
                            "bounding_box": {
                                "top_left_x": 10,
                                "top_left_y": 20,
                                "bottom_right_x": 200,
                                "bottom_right_y": 100,
                            },
                            "bounding_box_normalized": {
                                "top_left_x": 0.1,
                                "top_left_y": 0.2,
                                "bottom_right_x": 0.8,
                                "bottom_right_y": 0.6,
                            },
                        },
                    },
                ],
            }
        ],
        "meta": {"billed_units": {"pages": 1}},
    }


def test_parse_page_calls_staging_and_shapes_output() -> None:
    session = FakeSession(
        [
            FakeResponse(
                200,
                parse_response(),
                {"x-request-id": "request-123"},
            )
        ]
    )
    client = ParserClient(config(), session=session, sleeper=lambda _: None)

    result = client.parse_page(b"page")

    assert session.calls[0]["url"] == config().api_url
    assert session.calls[0]["timeout"] == 12
    assert session.calls[0]["headers"]["Authorization"] == "Bearer secret-key"
    assert session.calls[0]["json"]["output_format"] == "blocks"
    assert result["text_output"].startswith("# Invoice")
    assert '"id": "parse-response-123"' in result["raw_output"]
    assert result["boxes"][0]["label"] == "table"
    assert result["usage"] == {"billed_units": {"pages": 1}}
    assert result["request_id"] == "request-123"
    assert result["model"] == config().model


def test_retries_retryable_responses() -> None:
    session = FakeSession(
        [
            FakeResponse(503, {}, {"Retry-After": "0"}),
            FakeResponse(200, parse_response("ok")),
        ]
    )
    delays = []
    client = ParserClient(config(), session=session, sleeper=delays.append)

    result = client.parse_page(b"page")

    assert result["text_output"].startswith("ok")
    assert len(session.calls) == 2
    assert delays == [0.0]


def test_rejects_invalid_success_response() -> None:
    session = FakeSession([FakeResponse(200, {"id": "missing-pages"})])
    client = ParserClient(config(), session=session, sleeper=lambda _: None)

    with pytest.raises(ParserServiceError, match="invalid response"):
        client.parse_page(b"page")


def test_retries_transport_errors_then_returns_safe_error() -> None:
    session = FakeSession(
        [
            requests.ConnectionError("private network details"),
            requests.Timeout("private network details"),
            requests.ConnectionError("private network details"),
        ]
    )
    client = ParserClient(config(), session=session, sleeper=lambda _: None)

    with pytest.raises(ParserServiceError, match="could not be reached"):
        client.parse_page(b"page")


def test_rejects_non_retryable_response_without_leaking_body() -> None:
    session = FakeSession([FakeResponse(401, {"message": "secret detail"})])
    client = ParserClient(config(), session=session, sleeper=lambda _: None)

    with pytest.raises(ParserServiceError) as error:
        client.parse_page(b"page")

    assert "401" in str(error.value)
    assert "secret detail" not in str(error.value)