391 lines
14 KiB
Python
391 lines
14 KiB
Python
"""OcrClient 治理循环测试(M3 设计 §5): 换源/熔断口径/stall/取消/G1 契约。
|
|
|
|
循环与 EmbeddingClient 同构(设计 §2.A 有限重复裁决);本文件的
|
|
retry_exhausted/circuit_open/stalled 三组断言即设计 §6 ③ 的 G1 契约钉。
|
|
"""
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
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 TestCancellation:
|
|
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 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"] is not None
|
|
assert recorder.rows[1]["error"] is None
|
|
assert recorder.rows[1]["prompt_tokens"] == 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()
|