Files
PolyGateway/tests/unit/test_ocr_client.py
T

466 lines
17 KiB
Python

"""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 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()
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 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()