PolyGateway/tests/unit/test_ocr_client.py

"""OcrClient 治理循环测试(M3 设计 §5): 换源/熔断口径/stall/取消/G1 契约。

循环与 EmbeddingClient 同构(设计 §2.A 有限重复裁决);本文件的
retry_exhausted/circuit_open/stalled 三组断言即设计 §6 ③ 的 G1 契约钉。
"""

import asyncio

import pytest
from loguru import logger

from polygateway.backends.memory.breaker import InMemoryGate
from polygateway.backends.memory.limiter import InMemoryLimiter
from polygateway.errors import (
    AllSourcesExhausted,
    CircuitOpenError,
    RequestRejectedError,
    ResultInvalidError,
    SourceDeadError,
    TransientError,
)
from polygateway.ocr import OcrClient
from polygateway.types import (
    BackpressurePolicy,
    BreakerConfig,
    GlobalLimits,
    OcrLayoutElement,
    OcrLayoutTransportResult,
    OcrTextTransportResult,
    RetryPolicy,
    SourceConfig,
)

_BREAKER = BreakerConfig(fail_threshold=3, cooldown_s=60.0, probe_ttl_s=120.0)
_NO_GLOBAL = GlobalLimits(max_concurrency=0, rpm=0, tpm=0)
_TEXT_OK = OcrTextTransportResult(text="LINE-1", raw={"task_type": "text"})
_LAYOUT_OK = OcrLayoutTransportResult(
    elements=[OcrLayoutElement(type="table", bbox=(41.0, 48.0, 218.0, 282.0), page_index=0)],
    page_sizes=[(759.0, 540.0)],
    raw={"success": True},
)


def _src(**overrides):
    base = {
        "name": "m1",
        "provider": "monkey",
        "base_url": "http://10.77.0.20:7866",
        "api_key": "none",
        "model": "monkey-ocr",
        "timeout_s": 120.0,
    }
    base.update(overrides)
    return SourceConfig(**base)


class ScriptedOcrTransport:
    """按脚本响应: Exception / "text" / "layout" / "hang";两方法共用一份脚本。"""

    def __init__(self, script):
        self.script = list(script)
        self.calls = []

    async def _next(self, method, source, call_id):
        self.calls.append((method, source.name, call_id))
        action = self.script.pop(0)
        if isinstance(action, Exception):
            raise action
        if action == "hang":
            await asyncio.Event().wait()
        return _TEXT_OK if action == "text" else _LAYOUT_OK

    async def recognize_text(self, *, image, source, call_id):
        return await self._next("text", source, call_id)

    async def parse_layout(self, *, image, source, call_id):
        return await self._next("layout", source, call_id)

    async def check_health(self, *, source):
        raise NotImplementedError


class ClockAdvancingOcrTransport(ScriptedOcrTransport):
    """按脚本 [(推进秒数, 动作), ...] 在一次尝试内部推进时钟(issue #8)。

    stall 口径要区分"时间花在哪",故必须能让时钟只在 transport 内前进。
    """

    def __init__(self, script, clock):
        super().__init__([a for _, a in script])
        self._advances = [d for d, _ in script]
        self.clock = clock

    async def _next(self, method, source, call_id):
        self.clock.advance(self._advances.pop(0))
        return await super()._next(method, source, call_id)


class StaticSelector:
    def order(self, sources, stats):
        return list(sources)


class RecordingSelector(StaticSelector):
    def __init__(self):
        self.outcomes = []

    def record_outcome(self, source_name, ok):
        self.outcomes.append((source_name, ok))

    def health(self, source_name):
        return 1.0


class RecordingGate:
    """InMemoryGate 包装: 录 record_success 的 count_attempt 与 release_probe。"""

    def __init__(self, inner):
        self._inner = inner
        self.successes = []
        self.probe_releases = 0

    async def try_enter(self, source_name, owner):
        return await self._inner.try_enter(source_name, owner)

    async def record_success(self, entry, *, count_attempt=True):
        self.successes.append((entry.source_name, count_attempt))
        return await self._inner.record_success(entry, count_attempt=count_attempt)

    async def record_failure(self, entry, reason, force_open):
        return await self._inner.record_failure(entry, reason, force_open)

    async def release_probe(self, entry):
        self.probe_releases += 1
        return await self._inner.release_probe(entry)

    async def retry_after_s(self, sources):
        return await self._inner.retry_after_s(sources)


class _MemoryRecorder:
    def __init__(self):
        self.rows = []

    async def record_llm_call(self, **fields):
        self.rows.append(fields)


def _client(sources, script, **overrides):
    limiter = InMemoryLimiter(
        scope="ocr",
        sources={s.name: s for s in sources},
        global_limits=_NO_GLOBAL,
        lease_ttl_s=100.0,
    )
    gate = RecordingGate(InMemoryGate(config=_BREAKER))
    kwargs = {
        "scope": "ocr",
        "sources": sources,
        "selector": RecordingSelector(),
        "limiter": limiter,
        "breaker": gate,
        "transport": ScriptedOcrTransport(script),
        "retry": RetryPolicy(max_attempts=3, backoff_base_s=0.001, backoff_max_s=0.01),
        "backpressure": BackpressurePolicy(stall_window_s=300.0, poll_interval_s=0.001),
    }
    kwargs.update(overrides)
    return OcrClient(**kwargs), limiter, gate


class TestSuccessPaths:
    async def test_text_result_fields(self):
        client, _, gate = _client([_src()], ["text"])
        r = await client.recognize_text(b"jpg")
        assert r.text == "LINE-1"
        assert r.source_name == "m1"
        assert r.usage.prompt_tokens == 0 and r.usage.completion_tokens == 0
        assert r.latency_ms >= 0 and r.call_id and r.raw["task_type"] == "text"
        assert gate.successes == [("m1", True)]
        assert client._selector.outcomes == [("m1", True)]

    async def test_layout_result_fields(self):
        client, _, _ = _client([_src()], ["layout"])
        r = await client.parse_layout(b"jpg")
        assert r.elements[0].type == "table" and r.page_sizes == [(759.0, 540.0)]

    async def test_input_validation(self):
        client, _, _ = _client([_src()], [])
        with pytest.raises(TypeError):
            await client.recognize_text("not-bytes")
        with pytest.raises(ValueError):
            await client.parse_layout(b"")


class TestFailover:
    async def test_transient_retries_with_backoff(self):
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        client, _, _ = _client(
            [_src()], [TransientError("boom", status_code=500), "text"], sleep=fake_sleep
        )
        r = await client.recognize_text(b"jpg")
        assert r.text == "LINE-1"
        assert any(s > 0 for s in sleeps)  # Transient 退避后重试
        assert client._selector.outcomes == [("m1", False), ("m1", True)]

    async def test_source_dead_switches_immediately(self):
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        s1, s2 = _src(name="m1"), _src(name="m2", base_url="http://10.77.0.20:7867")
        client, _, _ = _client(
            [s1, s2], [SourceDeadError("401", status_code=401), "text"], sleep=fake_sleep
        )
        r = await client.recognize_text(b"jpg")
        assert r.source_name == "m2"
        assert not [s for s in sleeps if s > 0.0005]  # dead 立即换源不退避

    async def test_retry_exhausted_carries_g1_fields(self):
        client, _, _ = _client(
            [_src()],
            [TransientError("1", status_code=500)] * 3,
        )
        with pytest.raises(AllSourcesExhausted) as ei:
            await client.recognize_text(b"jpg")
        exc = ei.value
        assert exc.reason == "retry_exhausted"
        assert exc.per_source_reasons == {"m1": "network_error"}  # G1: 逐源原因
        assert exc.retry_after_s > 0  # G1: 可延期重投

    async def test_circuit_open_carries_g1_fields(self):
        client, _, _ = _client([_src()], [SourceDeadError("401", status_code=401)])
        with pytest.raises(CircuitOpenError) as ei:
            await client.recognize_text(b"jpg")
        exc = ei.value
        assert exc.per_source_reasons["m1"] == "circuit_open"
        assert exc.retry_after_s > 0


class TestTerminalOutcomes:
    async def test_result_invalid_counts_no_attempt_and_no_feed(self):
        client, _, gate = _client([_src()], [ResultInvalidError("bad zip")])
        with pytest.raises(ResultInvalidError):
            await client.parse_layout(b"jpg")
        assert gate.successes == [("m1", False)]  # 记成功但不入失败率窗
        assert client._selector.outcomes == []  # 坏结果 ≠ 坏服务,不喂健康

    async def test_rejected_with_status_counts_no_attempt(self):
        client, _, gate = _client([_src()], [RequestRejectedError("parse failed", status_code=200)])
        with pytest.raises(RequestRejectedError):
            await client.parse_layout(b"jpg")
        assert gate.successes == [("m1", False)]

    async def test_rejected_without_status_no_gate_success(self):
        client, _, gate = _client([_src()], [RequestRejectedError("local refuse")])
        with pytest.raises(RequestRejectedError):
            await client.recognize_text(b"jpg")
        assert gate.successes == []  # 服务未响应: 不记成功(非探针也不归还)


class TestBackpressure:
    async def test_fail_fast_when_quota_full(self):
        src = _src(max_concurrency=1)
        client, limiter, _ = _client([src], ["text"], quota_full="fail_fast")
        permit = await limiter.acquire("m1", 0)  # 占满唯一并发位
        try:
            with pytest.raises(AllSourcesExhausted) as ei:
                await client.recognize_text(b"jpg")
            assert ei.value.reason == "quota_exhausted"
        finally:
            await permit.settle(0)
            await permit.release()

    async def test_wait_until_permit_freed(self):
        src = _src(max_concurrency=1)
        client, limiter, _ = _client([src], ["text"])
        permit = await limiter.acquire("m1", 0)

        async def free_later():
            await asyncio.sleep(0.01)
            await permit.settle(0)
            await permit.release()

        release_task = asyncio.create_task(free_later())
        r = await client.recognize_text(b"jpg")
        await release_task
        assert r.text == "LINE-1"  # quota_full=wait: 等到许可释放而非报错

    async def test_stalled_when_no_progress(self):
        src = _src(max_concurrency=1)
        client, limiter, _ = _client(
            [src],
            ["text"],
            backpressure=BackpressurePolicy(stall_window_s=0.01, poll_interval_s=0.001),
        )
        permit = await limiter.acquire("m1", 0)  # 永不释放且全局无进展
        try:
            with pytest.raises(AllSourcesExhausted) as ei:
                await client.recognize_text(b"jpg")
            assert ei.value.reason == "stalled"
            assert ei.value.per_source_reasons["m1"] == "rate_limited"
        finally:
            await permit.settle(0)
            await permit.release()

    async def test_single_timeout_does_not_exhaust_stall_budget(self):
        """issue #8: 一次耗满 timeout 的尝试不得吃掉 stall 预算而使重试失效。

        OCR 只有 `_on_no_runnable` 一处 stall 判定,故失效链条是"先超时一次
        (墙钟耗尽)→ 再遇到无可用源 → 判死"。此处正是这条路径。
        """
        clock = FakeClock()
        limiter = held = None
        rounds = []

        async def toggle_permit(_seconds):
            """首次退避占满 permit,迫使下一轮走 _on_no_runnable;之后放行。"""
            nonlocal held
            rounds.append(_seconds)
            if len(rounds) > 10:
                raise RuntimeError("超过 10 次轮询仍未判死/未获 permit")
            if len(rounds) == 1:
                held = await limiter.acquire("m1", 0)
            else:
                await held.settle(0)
                await held.release()

        transport = ClockAdvancingOcrTransport(
            [(300.1, TransientError("timeout", status_code=504)), (0.0, "text")], clock
        )
        client, limiter, _ = _client(
            [_src(max_concurrency=1)], [], now=clock, sleep=toggle_permit, transport=transport
        )
        r = await client.recognize_text(b"jpg")
        assert r.text == "LINE-1"
        assert len(transport.calls) == 2  # 第二次尝试确实发出了


class FakeClock:
    def __init__(self, start=1000.0):
        self.t = start

    def __call__(self):
        return self.t

    def advance(self, seconds):
        self.t += seconds


class TestCancellation:
    async def test_cancel_during_probe_returns_probe(self):
        # 探针取消归还(铁律分支;verifier I3): 开路→冷却过后半开探针→
        # transport 挂起中取消 → release_probe 必须发生,后续可再探
        clock = FakeClock()
        inner = InMemoryGate(config=_BREAKER, now=clock)  # 门与 client 共用注入钟
        gate = RecordingGate(inner)
        client, _, _ = _client([_src()], ["hang"], now=clock, breaker=gate)
        entry = await inner.try_enter("m1", "setup")
        await inner.record_failure(entry, "source_dead", True)  # force_open → OPEN
        clock.advance(61.0)  # 越过 cooldown_s=60 → HALF_OPEN
        task = asyncio.create_task(client.recognize_text(b"jpg"))
        await asyncio.sleep(0.05)
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        assert gate.probe_releases >= 1  # 探针已归还,不悬挂
        probe = await inner.try_enter("m1", "again")
        assert probe.allowed and probe.is_probe  # 可再探 = 未悬挂的行为证据

    async def test_cancel_during_backoff_sleep_propagates(self):
        async def cancelling_sleep(delay):
            raise asyncio.CancelledError()

        client, limiter, _ = _client(
            [_src()], [TransientError("boom", status_code=500)], sleep=cancelling_sleep
        )
        with pytest.raises(asyncio.CancelledError):
            await client.recognize_text(b"jpg")
        stats = await limiter.source_stats("m1")
        assert stats.inflight == 0  # 退避期取消: permit 已在 finally 释放

    async def test_cancel_during_transport_releases_permit(self):
        src = _src(max_concurrency=1)
        client, limiter, _ = _client([src], ["hang"])
        task = asyncio.create_task(client.recognize_text(b"jpg"))
        await asyncio.sleep(0.01)
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        stats = await limiter.source_stats("m1")
        assert stats.inflight == 0  # permit 在 finally 释放


class TestCheckHealth:
    class _HealthTransport(ScriptedOcrTransport):
        def __init__(self, mapping):
            super().__init__([])
            self.mapping = mapping

        async def check_health(self, *, source):
            result = self.mapping[source.name]
            if result == "hang":
                await asyncio.Event().wait()
            if isinstance(result, Exception):
                raise result
            return result

    async def test_per_source_dict(self):
        s1, s2 = _src(name="m1"), _src(name="m2", base_url="http://10.77.0.20:7867")
        client, _, _ = _client([s1, s2], [])
        client._transport = self._HealthTransport({"m1": True, "m2": False})
        assert await client.check_health() == {"m1": True, "m2": False}

    async def test_cancellation_passes_through(self):
        client, _, _ = _client([_src()], [])
        client._transport = self._HealthTransport({"m1": "hang"})
        task = asyncio.create_task(client.check_health())
        await asyncio.sleep(0.01)
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task


class TestExtraBodyStripped:
    """issue #4 决策 G: OCR 路径只发 multipart 表单,剥离 extra_body 并 warning。"""

    def test_stripped_with_warning_but_assembly_succeeds(self):
        messages: list[str] = []
        sink_id = logger.add(messages.append, level="WARNING")
        try:
            client, _, _ = _client([_src(extra_body={"temperature": 0})], ["text"])
        finally:
            logger.remove(sink_id)
        assert client._sources[0].extra_body == {}
        assert any("extra_body" in m for m in messages)

    async def test_telemetry_never_records_a_parameter_that_was_not_sent(self):
        """不剥离则审计表会显示这次 OCR 带了 temperature=0——数据造假。"""
        recorder = _MemoryRecorder()
        client, _, _ = _client([_src(extra_body={"temperature": 0})], ["text"], telemetry=recorder)
        await client.recognize_text(b"IMG")
        assert recorder.rows[0]["sampling"] is None


class TestTelemetry:
    async def test_success_and_failure_recorded_without_image_bytes(self):
        recorder = _MemoryRecorder()
        client, _, _ = _client(
            [_src()],
            [TransientError("boom", status_code=500), "text"],
            telemetry=recorder,
        )
        await client.recognize_text(b"RAW-IMAGE-BYTES")
        assert len(recorder.rows) == 2  # 失败尝试与成功尝试均必录
        for row in recorder.rows:
            assert "<ocr:text image_bytes=15>" in row["messages"]
            assert "RAW-IMAGE-BYTES" not in row["messages"]  # 图像字节绝不入库
        assert recorder.rows[0]["error"].startswith("TransientError:")  # 类名前缀口径
        assert recorder.rows[1]["error"] is None
        assert recorder.rows[1]["prompt_tokens"] == 0

    async def test_success_row_stays_measured_and_settles_zero(self):
        """OCR 的 0 token 是**事实**而非未知(est_tokens 解耦设计 §3.3 剔出决定)。

        三态化不得把 OCR 成功行改成 `unavailable`——那会灌水缺口度量
        `COUNT(*) WHERE usage_source='unavailable'`;settle 恒 0 的差异①同样不动。
        """
        recorder = _MemoryRecorder()
        client, limiter, _ = _client([_src(tpm=1000, est_tokens=400)], ["text"], telemetry=recorder)
        await client.recognize_text(b"jpg")
        row = recorder.rows[0]
        assert row["usage_source"] == "measured"
        assert row["prompt_tokens"] == 0 and row["completion_tokens"] == 0
        assert row["error"] is None
        assert (await limiter.source_stats("m1")).tpm_used == 0  # settle(0) 全额退回预扣


class TestOcrCallerDimensions:
    """issue #11: 调用方自定义维度必须沿 OCR 链四层透传到每一行遥测。

    OCR 行与 chat 行落在同一张 `llm_calls` 表: 不覆盖这条链会让同一张表里
    一部分行有租户归属、一部分永远空白,而"先启用后加列则归属无法还原"。
    """

    async def test_recognize_text_row_carries_dimensions(self):
        recorder = _MemoryRecorder()
        client, _, _ = _client([_src()], ["text"], telemetry=recorder)
        await client.recognize_text(b"jpg", tenant_id="t1", meta={"batch": "b-42"})
        assert recorder.rows[0]["tenant_id"] == "t1"
        assert recorder.rows[0]["meta"] == '{"batch": "b-42"}'

    async def test_parse_layout_row_carries_dimensions(self):
        """两个公共方法都是入口: 只测一个会漏掉另一个的透传缺口。"""
        recorder = _MemoryRecorder()
        client, _, _ = _client([_src()], ["layout"], telemetry=recorder)
        await client.parse_layout(b"jpg", tenant_id="t2", meta={"batch": "b-43"})
        assert recorder.rows[0]["tenant_id"] == "t2"
        assert recorder.rows[0]["meta"] == '{"batch": "b-43"}'

    async def test_failed_attempt_row_also_carries_dimensions(self):
        """失败行同样需要归属: 某租户的请求没被服务,正是审计最需要的一行。"""
        recorder = _MemoryRecorder()
        client, _, _ = _client(
            [_src()],
            [TransientError("boom", status_code=500), "text"],
            telemetry=recorder,
        )
        await client.recognize_text(b"jpg", tenant_id="t1", meta={"batch": "b-42"})
        assert len(recorder.rows) == 2  # 失败尝试 + 成功尝试
        assert [r["tenant_id"] for r in recorder.rows] == ["t1", "t1"]
        assert [r["meta"] for r in recorder.rows] == ['{"batch": "b-42"}'] * 2

    @pytest.mark.parametrize("method", ["recognize_text", "parse_layout"])
    async def test_invalid_meta_rejected_before_any_telemetry(self, method):
        """校验必须早于遥测: 链路内的失败都被降级成 warning,放下游等于没有校验。"""
        recorder = _MemoryRecorder()
        client, _, _ = _client([_src()], ["text"], telemetry=recorder)
        with pytest.raises(ValueError, match="meta"):
            await getattr(client, method)(b"jpg", meta={"Bad Key": 1})
        assert recorder.rows == []
        assert client._transport.calls == []  # 连调用都没发出

    async def test_defaults_land_as_sentinels(self):
        recorder = _MemoryRecorder()
        client, _, _ = _client([_src()], ["text"], telemetry=recorder)
        await client.recognize_text(b"jpg")
        assert recorder.rows[0]["tenant_id"] == ""  # 空串哨兵,不是 None
        assert recorder.rows[0]["meta"] == "{}"


class TestAssembly:
    _ENV = {
        "OCR__MONKEY__1__BASE_URL": "http://10.77.0.20:7866",
        "OCR__MONKEY__1__API_KEY": "none",
        "OCR__MONKEY__1__MODEL": "monkey-ocr",
        "OCR__MONKEY__1__TIMEOUT_S": "120",
        "LLM_MAX_RETRIES": "3",
        "LLM_RETRY_BASE_DELAY": "2.0",
        "LLM_RETRY_MAX_DELAY": "30.0",
        "LLM_CIRCUIT_BREAKER_THRESHOLD": "5",
        "LLM_CIRCUIT_BREAKER_COOLDOWN": "60",
        "PGW_CACHE_BACKEND": "none",
        "PGW_TELEMETRY_BACKEND": "none",
    }

    async def test_from_env_assembles(self):
        client = OcrClient.from_env("OCR", env=dict(self._ENV))
        try:
            assert client._scope == "ocr"
        finally:
            await client.aclose()

    async def test_non_monkey_provider_rejected(self):
        env = {k.replace("MONKEY", "GLM"): v for k, v in self._ENV.items()}
        with pytest.raises(ValueError, match="monkey"):
            OcrClient.from_env("OCR", env=env)

    async def test_aclose_idempotent(self):
        client = OcrClient.from_env("OCR", env=dict(self._ENV))
        await client.aclose()
        await client.aclose()