fix(provider_keys): 提升 Codex 配额异步同步的可靠性与可观测性

- 为 flush 引入 FlushResult,统一返回更新数与重试批次
- 增加指数退避与失败日志限流,成功后重置 backoff
- 批量提交失败时回退到单条提交,降低整批失败风险
- 补充测试,覆盖 flush 重试与提交失败回退场景
This commit is contained in:
AAEE86
2026-02-28 16:54:37 +08:00
parent 92e9caf57e
commit 6a8b5e6c8e
2 changed files with 238 additions and 33 deletions

View File

@@ -8,6 +8,45 @@ import pytest
from src.services.provider_keys import codex_quota_sync_dispatcher as dispatcher_module
def _flush_ok(batch: dict[str, dict[str, Any]]) -> dispatcher_module.FlushResult:
return dispatcher_module.FlushResult(
queued_count=len(batch),
updated_count=len(batch),
retry_batch={},
)
class _FakeNestedTxn:
def __enter__(self) -> "_FakeNestedTxn":
return self
def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
_ = exc_type, exc, tb
return None
class _FakeSession:
def __init__(self, *, fail_commit: bool = False) -> None:
self.fail_commit = fail_commit
self.commit_calls = 0
self.rollback_calls = 0
self.closed = False
def begin_nested(self) -> _FakeNestedTxn:
return _FakeNestedTxn()
def commit(self) -> None:
self.commit_calls += 1
if self.fail_commit and self.commit_calls == 1:
raise RuntimeError("commit failed once")
def rollback(self) -> None:
self.rollback_calls += 1
def close(self) -> None:
self.closed = True
def test_dispatch_falls_back_to_sync_when_dispatcher_not_running(
monkeypatch: pytest.MonkeyPatch,
) -> None:
@@ -46,8 +85,9 @@ async def test_dispatcher_deduplicates_headers_by_provider_api_key_id() -> None:
dispatcher = dispatcher_module.CodexQuotaSyncDispatcher(flush_interval_seconds=0.01)
flushed_batches: list[dict[str, dict[str, Any]]] = []
def _fake_flush(batch: dict[str, dict[str, Any]]) -> None:
def _fake_flush(batch: dict[str, dict[str, Any]]) -> dispatcher_module.FlushResult:
flushed_batches.append(batch)
return _flush_ok(batch)
dispatcher._flush_batch_sync = _fake_flush # type: ignore[method-assign]
@@ -83,12 +123,13 @@ async def test_dispatcher_retries_batch_after_flush_error() -> None:
flushed_batches: list[dict[str, dict[str, Any]]] = []
flush_attempts = 0
def _flaky_flush(batch: dict[str, dict[str, Any]]) -> None:
def _flaky_flush(batch: dict[str, dict[str, Any]]) -> dispatcher_module.FlushResult:
nonlocal flush_attempts
flush_attempts += 1
if flush_attempts == 1:
raise RuntimeError("temporary flush error")
flushed_batches.append(batch)
return _flush_ok(batch)
dispatcher._flush_batch_sync = _flaky_flush # type: ignore[method-assign]
@@ -110,14 +151,54 @@ async def test_dispatcher_retries_batch_after_flush_error() -> None:
assert merged["retry-key"] == {"x-codex-primary-used-percent": "7"}
def test_flush_batch_falls_back_to_single_item_commit(monkeypatch: pytest.MonkeyPatch) -> None:
dispatcher = dispatcher_module.CodexQuotaSyncDispatcher(flush_interval_seconds=0.01)
main_session = _FakeSession(fail_commit=True)
fallback_session_a = _FakeSession()
fallback_session_b = _FakeSession()
sessions: list[_FakeSession] = [main_session, fallback_session_a, fallback_session_b]
def _fake_create_session() -> _FakeSession:
return sessions.pop(0)
def _fake_sync(
*,
db: Any,
provider_api_key_id: str | None,
response_headers: dict[str, Any] | None,
) -> bool:
_ = db, provider_api_key_id, response_headers
return True
monkeypatch.setattr(dispatcher_module, "create_session", _fake_create_session)
monkeypatch.setattr(dispatcher_module, "sync_codex_quota_from_response_headers", _fake_sync)
result = dispatcher._flush_batch_sync(
{
"key-a": {"x-codex-primary-used-percent": "1"},
"key-b": {"x-codex-primary-used-percent": "2"},
}
)
assert result.updated_count == 2
assert result.retry_batch == {}
assert main_session.rollback_calls == 1
assert fallback_session_a.commit_calls == 1
assert fallback_session_b.commit_calls == 1
assert main_session.closed is True
assert fallback_session_a.closed is True
assert fallback_session_b.closed is True
@pytest.mark.asyncio
async def test_dispatch_uses_async_queue_when_dispatcher_running() -> None:
dispatcher = dispatcher_module.CodexQuotaSyncDispatcher(flush_interval_seconds=0.01)
flushed_batches: list[dict[str, dict[str, Any]]] = []
dispatcher_module._dispatcher_instance = dispatcher
def _fake_flush(batch: dict[str, dict[str, Any]]) -> None:
def _fake_flush(batch: dict[str, dict[str, Any]]) -> dispatcher_module.FlushResult:
flushed_batches.append(batch)
return _flush_ok(batch)
dispatcher._flush_batch_sync = _fake_flush # type: ignore[method-assign]