i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: Antigravity/Codex 服务重构为插件化适配器架构

Browse Source 
- 将 Antigravity 和 Codex 从独立模块迁移至 src/services/provider/adapters/ 插件体系
- 新增 provider_types 和 oauth_token 模块，移除 maintenance_scheduler 中的 OAuth 定时刷新
- 增强 admin API：扩展 keys 和 provider_query 端点，新增 dashboard 路由
- 大幅增强 ProviderDetailDrawer 组件，新增 AntigravityQuotaDialog
- 改进 handler 基类（chat/cli）和错误分类器
- 优化 fetch_scheduler 和 upstream_fetcher
- 前端 UI 组件清理和优化
- 更新测试以匹配新模块结构
This commit is contained in:
fawney19
2026-02-06 16:37:06 +08:00
parent e01dfee41d
commit b88fb6273b
112 changed files with 5193 additions and 1992 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
tests/api/handlers/base/test_antigravity_v1internal.py
View File

@@ -6,19 +6,19 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from src.services.antigravity.envelope import (
from src.api.handlers.base.stream_context import StreamContext
from src.api.handlers.gemini_cli.handler import GeminiCliMessageHandler
from src.services.provider.adapters.antigravity.envelope import (
unwrap_v1internal_response,
wrap_v1internal_request,
)
from src.api.handlers.base.stream_context import StreamContext
from src.api.handlers.gemini_cli.handler import GeminiCliMessageHandler
def _make_handler() -> GeminiCliMessageHandler:
return GeminiCliMessageHandler(
db=MagicMock(),
user=SimpleNamespace(id=1),
api_key=SimpleNamespace(id=1),
db=MagicMock(), # type: ignore[arg-type]
user=SimpleNamespace(id=1), # type: ignore[arg-type]
api_key=SimpleNamespace(id=1), # type: ignore[arg-type]
request_id="req_1",
client_ip="127.0.0.1",
user_agent="pytest",
@@ -73,7 +73,9 @@ def test_convert_sse_line_unwraps_before_convert_stream_chunk() -> None:
seen: dict[str, object] = {}
class _DummyRegistry:
def convert_stream_chunk(self, data_obj, *_args, **_kwargs): # noqa: ANN001
def convert_stream_chunk(
self, data_obj: object, *_args: object, **_kwargs: object
) -> list[str]:
seen["data_obj"] = data_obj
return []
@@ -84,8 +86,8 @@ def test_convert_sse_line_unwraps_before_convert_stream_chunk() -> None:
_lines, _events = handler._convert_sse_line(ctx, v1_line, [])
assert isinstance(seen.get("data_obj"), dict)
assert "response" not in seen["data_obj"] # 已解包
assert "_v1internal_response_id" in seen["data_obj"]
assert "response" not in seen["data_obj"] # type: ignore[operator]
assert "_v1internal_response_id" in seen["data_obj"] # type: ignore[operator]
def test_handle_sse_event_unwraps_for_antigravity() -> None:
@@ -118,14 +120,17 @@ def test_handle_sse_event_unwraps_for_antigravity() -> None:
def test_handle_sse_event_caches_thought_signature_for_antigravity() -> None:
from src.services.antigravity.signature_cache import signature_cache
from src.services.provider.adapters.antigravity.signature_cache import signature_cache
signature_cache._cache.clear() # type: ignore[attr-defined]
signature_cache.clear()
handler = _make_handler()
ctx = StreamContext(model="claude-sonnet-4-5", api_format="gemini:cli")
ctx.provider_type = "antigravity"
# 签名须 >= MIN_SIGNATURE_LENGTH(50)，否则会被忽略
long_sig = "a" * 60
payload = {
"candidates": [
{
@@ -134,7 +139,7 @@ def test_handle_sse_event_caches_thought_signature_for_antigravity() -> None:
{
"text": "t1",
"thought": True,
"thoughtSignature": "sig-abc",
"thoughtSignature": long_sig,
}
]
}
@@ -145,7 +150,7 @@ def test_handle_sse_event_caches_thought_signature_for_antigravity() -> None:
with patch.object(handler, "_process_event_data") as _mock_process:
handler._handle_sse_event(ctx, None, json.dumps(payload), record_chunk=False)
assert signature_cache.get_or_dummy("claude-sonnet-4-5", "t1") == "sig-abc"
assert signature_cache.get_or_dummy("claude-sonnet-4-5", "t1") == long_sig
def test_provider_type_drives_antigravity_usage_path() -> None:
@@ -167,13 +172,13 @@ async def test_antigravity_forces_conversion_path_in_stream_with_prefetch() -> N
ctx.needs_conversion = False # same-format case normally would passthrough
class _AsyncIter:
def __init__(self, items): # noqa: ANN001
def __init__(self, items: list[bytes]) -> None:
self._it = iter(items)
def __aiter__(self):
def __aiter__(self) -> _AsyncIter:
return self
async def __anext__(self):
async def __anext__(self) -> bytes:
try:
return next(self._it)
except StopIteration as e:
@@ -193,7 +198,7 @@ async def test_antigravity_forces_conversion_path_in_stream_with_prefetch() -> N
) as mock_convert:
out = []
async for chunk in handler._create_response_stream_with_prefetch(
ctx, byte_iter, response_ctx, http_client, prefetched
ctx, byte_iter, response_ctx, http_client, prefetched # type: ignore[arg-type]
):
out.append(chunk)

10
tests/core/test_provider_oauth_utils_antigravity.py
View File

@@ -13,9 +13,12 @@ async def test_enrich_auth_config_antigravity_adds_project_id_and_email() -> Non
token_response: dict[str, object] = {}
with (
patch("src.core.provider_oauth_utils.fetch_google_email", AsyncMock(return_value="u@example.com")),
patch(
"src.services.antigravity.client.load_code_assist",
"src.core.provider_oauth_utils.fetch_google_email",
AsyncMock(return_value="u@example.com"),
),
patch(
"src.services.provider.adapters.antigravity.client.load_code_assist",
AsyncMock(
return_value={
"cloudaicompanionProject": "project-1",
@@ -34,5 +37,4 @@ async def test_enrich_auth_config_antigravity_adds_project_id_and_email() -> Non
assert out["email"] == "u@example.com"
assert out["project_id"] == "project-1"
assert out["tier"] == "PAID"
assert out["tier"] == "Pro"

49
tests/e2e/test_video_cancel_e2e.py
View File

@@ -41,8 +41,12 @@ async def test_video_cancel_openai_route_end_to_end(monkeypatch: pytest.MonkeyPa
# Pipeline auth/quota/audit shortcuts
user = SimpleNamespace(id="u1", username="u1", role="user", quota_usd=None, used_usd=0.0)
api_key = SimpleNamespace(id="ak1", user_id="u1", is_standalone=False)
monkeypatch.setattr(pipeline.auth_service, "authenticate_api_key", lambda _db, _k: (user, api_key))
monkeypatch.setattr(pipeline.usage_service, "check_user_quota", lambda *_args, **_kwargs: (True, "ok"))
monkeypatch.setattr(
pipeline.auth_service, "authenticate_api_key", lambda _db, _k: (user, api_key)
)
monkeypatch.setattr(
pipeline.usage_service, "check_user_quota", lambda *_args, **_kwargs: (True, "ok")
)
monkeypatch.setattr(pipeline.audit_service, "log_event", MagicMock())
# DB stubs used by TaskService.cancel
@@ -81,7 +85,7 @@ async def test_video_cancel_openai_route_end_to_end(monkeypatch: pytest.MonkeyPa
db = MagicMock()
def _query(model): # noqa: ANN001
def _query(model: type) -> MagicMock:
name = getattr(model, "__name__", "")
if name == "VideoTask":
return q_task
@@ -100,19 +104,28 @@ async def test_video_cancel_openai_route_end_to_end(monkeypatch: pytest.MonkeyPa
)
with (
patch("src.clients.http_client.HTTPClientPool.get_default_client_async", AsyncMock(return_value=upstream)),
patch(
"src.clients.http_client.HTTPClientPool.get_default_client_async",
AsyncMock(return_value=upstream),
),
patch("src.core.crypto.crypto_service.decrypt", lambda _v: "upstream-key"),
patch(
"src.services.provider.transport.build_provider_url",
lambda _endpoint, **_kwargs: "https://upstream.example.com/v1/videos",
),
patch("src.services.usage.service.UsageService.finalize_void", MagicMock(return_value=True)),
patch(
"src.services.usage.service.UsageService.finalize_void", MagicMock(return_value=True)
),
patch("src.services.usage.service.UsageService.void_settled", MagicMock()),
):
request = _make_request(
method="POST",
path="/v1/videos/t1/cancel",
headers={"authorization": "Bearer sk-test", "x-real-ip": "127.0.0.1", "user-agent": "pytest"},
headers={
"authorization": "Bearer sk-test",
"x-real-ip": "127.0.0.1",
"user-agent": "pytest",
},
body=b"",
)
adapter = OpenAIVideoAdapter()
@@ -142,8 +155,12 @@ async def test_video_cancel_gemini_route_end_to_end(monkeypatch: pytest.MonkeyPa
# Pipeline auth/quota/audit shortcuts
user = SimpleNamespace(id="u1", username="u1", role="user", quota_usd=None, used_usd=0.0)
api_key = SimpleNamespace(id="ak1", user_id="u1", is_standalone=False)
monkeypatch.setattr(pipeline.auth_service, "authenticate_api_key", lambda _db, _k: (user, api_key))
monkeypatch.setattr(pipeline.usage_service, "check_user_quota", lambda *_args, **_kwargs: (True, "ok"))
monkeypatch.setattr(
pipeline.auth_service, "authenticate_api_key", lambda _db, _k: (user, api_key)
)
monkeypatch.setattr(
pipeline.usage_service, "check_user_quota", lambda *_args, **_kwargs: (True, "ok")
)
monkeypatch.setattr(pipeline.audit_service, "log_event", MagicMock())
# DB stubs used by TaskService.cancel
@@ -185,7 +202,7 @@ async def test_video_cancel_gemini_route_end_to_end(monkeypatch: pytest.MonkeyPa
db = MagicMock()
_video_query_count = {"n": 0}
def _query(model): # noqa: ANN001
def _query(model: type) -> MagicMock:
name = getattr(model, "__name__", "")
if name == "VideoTask":
_video_query_count["n"] += 1
@@ -204,10 +221,17 @@ async def test_video_cancel_gemini_route_end_to_end(monkeypatch: pytest.MonkeyPa
)
with (
patch("src.clients.http_client.HTTPClientPool.get_default_client_async", AsyncMock(return_value=upstream)),
patch(
"src.clients.http_client.HTTPClientPool.get_default_client_async",
AsyncMock(return_value=upstream),
),
patch("src.core.crypto.crypto_service.decrypt", lambda _v: "upstream-key"),
patch("src.api.handlers.base.request_builder.get_provider_auth", AsyncMock(return_value=None)),
patch("src.services.usage.service.UsageService.finalize_void", MagicMock(return_value=True)),
patch(
"src.api.handlers.base.request_builder.get_provider_auth", AsyncMock(return_value=None)
),
patch(
"src.services.usage.service.UsageService.finalize_void", MagicMock(return_value=True)
),
patch("src.services.usage.service.UsageService.void_settled", MagicMock()),
):
request = _make_request(
@@ -230,4 +254,3 @@ async def test_video_cancel_gemini_route_end_to_end(monkeypatch: pytest.MonkeyPa
assert upstream.post.await_count == 1
assert upstream.post.call_args.args[0].endswith("/v1beta/operations/op123:cancel")
assert task.status == VideoStatus.CANCELLED.value

143
tests/services/antigravity/test_client.py
View File

@@ -6,31 +6,154 @@ from unittest.mock import AsyncMock, patch
import httpx
import pytest
from src.services.antigravity.client import load_code_assist
from src.services.antigravity.constants import DAILY_BASE_URL, PROD_BASE_URL
from src.services.provider.adapters.antigravity.client import (
fetch_available_models,
load_code_assist,
parse_retry_delay,
)
from src.services.provider.adapters.antigravity.constants import (
DAILY_BASE_URL,
PROD_BASE_URL,
SANDBOX_BASE_URL,
)
# ---------------------------------------------------------------------------
# load_code_assist
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_load_code_assist_falls_back_to_daily() -> None:
resp1 = httpx.Response(500, json={"error": {"message": "boom"}})
resp2 = httpx.Response(200, json={"cloudaicompanionProject": "project-1"})
async def test_load_code_assist_falls_back_on_500() -> None:
"""500 时 fallback 到下一个 URL（Sandbox → Daily → Prod 顺序）。"""
resp_fail = httpx.Response(500, json={"error": {"message": "boom"}})
resp_ok = httpx.Response(200, json={"cloudaicompanionProject": "project-1"})
client = SimpleNamespace(post=AsyncMock(side_effect=[resp1, resp2]))
client = SimpleNamespace(post=AsyncMock(side_effect=[resp_fail, resp_ok]))
with patch(
"src.clients.http_client.HTTPClientPool.get_proxy_client",
AsyncMock(return_value=client),
with (
patch(
"src.clients.http_client.HTTPClientPool.get_proxy_client",
AsyncMock(return_value=client),
),
patch(
"src.services.provider.adapters.antigravity.client.url_availability.get_ordered_urls",
return_value=[SANDBOX_BASE_URL, DAILY_BASE_URL, PROD_BASE_URL],
),
):
data = await load_code_assist("tok", proxy_config=None, timeout_seconds=1.0)
assert data["cloudaicompanionProject"] == "project-1"
assert client.post.await_count == 2
assert client.post.call_args_list[0].args[0] == f"{PROD_BASE_URL}/v1internal:loadCodeAssist"
assert client.post.call_args_list[0].args[0] == f"{SANDBOX_BASE_URL}/v1internal:loadCodeAssist"
assert client.post.call_args_list[1].args[0] == f"{DAILY_BASE_URL}/v1internal:loadCodeAssist"
@pytest.mark.asyncio
async def test_load_code_assist_4xx_does_not_fallback() -> None:
"""401/403 等 4xx 客户端错误不应 fallback，直接抛出。"""
resp_401 = httpx.Response(401, json={"error": "unauthorized"}, text="unauthorized")
client = SimpleNamespace(post=AsyncMock(return_value=resp_401))
with (
patch(
"src.clients.http_client.HTTPClientPool.get_proxy_client",
AsyncMock(return_value=client),
),
patch(
"src.services.provider.adapters.antigravity.client.url_availability.get_ordered_urls",
return_value=[SANDBOX_BASE_URL, DAILY_BASE_URL],
),
pytest.raises(RuntimeError, match="status=401"),
):
await load_code_assist("tok", proxy_config=None, timeout_seconds=1.0)
# 只调用了一次（没有 fallback 到第二个 URL）
assert client.post.await_count == 1
@pytest.mark.asyncio
async def test_load_code_assist_requires_token() -> None:
with pytest.raises(ValueError):
await load_code_assist("", proxy_config=None)
# ---------------------------------------------------------------------------
# fetch_available_models
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_fetch_available_models_falls_back_on_500() -> None:
resp_fail = httpx.Response(500, json={"error": {"message": "boom"}})
resp_ok = httpx.Response(
200,
json={
"models": {
"claude-sonnet-4": {
"displayName": "Claude Sonnet 4",
"quotaInfo": {"remainingFraction": 0.75, "resetTime": "2024-01-15T12:00:00Z"},
}
}
},
)
client = SimpleNamespace(post=AsyncMock(side_effect=[resp_fail, resp_ok]))
with (
patch(
"src.clients.http_client.HTTPClientPool.get_proxy_client",
AsyncMock(return_value=client),
),
patch(
"src.services.provider.adapters.antigravity.client.url_availability.get_ordered_urls",
return_value=[DAILY_BASE_URL, PROD_BASE_URL],
),
):
data = await fetch_available_models(
"tok",
project_id="project-1",
proxy_config=None,
timeout_seconds=1.0,
)
assert "models" in data
assert client.post.await_count == 2
assert (
client.post.call_args_list[0].args[0] == f"{DAILY_BASE_URL}/v1internal:fetchAvailableModels"
)
assert (
client.post.call_args_list[1].args[0] == f"{PROD_BASE_URL}/v1internal:fetchAvailableModels"
)
@pytest.mark.asyncio
async def test_fetch_available_models_requires_project_id() -> None:
with pytest.raises(ValueError):
await fetch_available_models("tok", project_id="", proxy_config=None)
# ---------------------------------------------------------------------------
# parse_retry_delay
# ---------------------------------------------------------------------------
def test_parse_retry_delay_from_retry_info() -> None:
error_json = '{"error": {"details": [{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "1.5s"}]}}'
delay = parse_retry_delay(error_json)
assert delay is not None
# 1500ms + 200ms buffer = 1700ms = 1.7s
assert 1.5 < delay < 2.0
def test_parse_retry_delay_from_quota_reset() -> None:
error_json = '{"error": {"details": [{"metadata": {"quotaResetDelay": "200ms"}}]}}'
delay = parse_retry_delay(error_json)
assert delay is not None
assert 0.3 < delay < 0.5
def test_parse_retry_delay_invalid() -> None:
assert parse_retry_delay("not json") is None
assert parse_retry_delay("{}") is None
assert parse_retry_delay('{"error": {}}') is None

132
tests/services/antigravity/test_signature_cache.py
View File

@@ -1,35 +1,133 @@
from __future__ import annotations
from src.services.antigravity.constants import DUMMY_THOUGHT_SIGNATURE
from src.services.antigravity.signature_cache import ThinkingSignatureCache
from src.services.provider.adapters.antigravity.constants import DUMMY_THOUGHT_SIGNATURE
from src.services.provider.adapters.antigravity.signature_cache import ThinkingSignatureCache
# 测试用签名（需 >= MIN_SIGNATURE_LENGTH=50）
_SIG_A = "a" * 60
_SIG_B = "b" * 60
_SIG_C = "c" * 60
_SIG_D = "d" * 60
_SIG_E = "e" * 60
def test_get_or_dummy_returns_dummy_for_gemini_models() -> None:
cache = ThinkingSignatureCache(maxsize=10)
cache = ThinkingSignatureCache()
assert cache.get_or_dummy("gemini-3-pro", "thinking...") == DUMMY_THOUGHT_SIGNATURE
def test_get_or_dummy_returns_none_for_non_gemini_models() -> None:
cache = ThinkingSignatureCache(maxsize=10)
cache = ThinkingSignatureCache()
assert cache.get_or_dummy("claude-sonnet", "thinking...") is None
def test_cached_signature_preferred() -> None:
cache = ThinkingSignatureCache(maxsize=10)
cache.cache("gemini-3-pro", "t", "sig-1")
assert cache.get_or_dummy("gemini-3-pro", "t") == "sig-1"
cache = ThinkingSignatureCache()
cache.cache("gemini-3-pro", "thinking-text", _SIG_A)
assert cache.get_or_dummy("gemini-3-pro", "thinking-text") == _SIG_A
def test_eviction_fifo() -> None:
cache = ThinkingSignatureCache(maxsize=4)
cache.cache("gemini-3-pro", "t1", "s1")
cache.cache("gemini-3-pro", "t2", "s2")
cache.cache("gemini-3-pro", "t3", "s3")
cache.cache("gemini-3-pro", "t4", "s4")
def test_short_signature_ignored() -> None:
"""短于 MIN_SIGNATURE_LENGTH 的签名不会被缓存。"""
cache = ThinkingSignatureCache()
cache.cache("gemini-3-pro", "text", "short")
# 未命中缓存，回退到 DUMMY
assert cache.get_or_dummy("gemini-3-pro", "text") == DUMMY_THOUGHT_SIGNATURE
# Trigger eviction (evict 1 key when maxsize=4)
cache.cache("gemini-3-pro", "t5", "s5")
# Oldest key should be gone
assert cache.get_or_dummy("gemini-3-pro", "t1") == DUMMY_THOUGHT_SIGNATURE
# ===== Layer 1: Tool Signatures =====
def test_tool_signature_cache() -> None:
cache = ThinkingSignatureCache()
cache.cache_tool_signature("toolu_123", _SIG_A)
assert cache.get_tool_signature("toolu_123") == _SIG_A
assert cache.get_tool_signature("toolu_999") is None
def test_tool_signature_short_ignored() -> None:
cache = ThinkingSignatureCache()
cache.cache_tool_signature("toolu_123", "short")
assert cache.get_tool_signature("toolu_123") is None
# ===== Layer 2: Thinking Families =====
def test_thinking_family_cache() -> None:
cache = ThinkingSignatureCache()
cache.cache_thinking_family(_SIG_A, "claude-3-5-sonnet")
assert cache.get_signature_family(_SIG_A) == "claude-3-5-sonnet"
assert cache.get_signature_family(_SIG_B) is None
# ===== Layer 3: Session Signatures =====
def test_session_signature_basic() -> None:
cache = ThinkingSignatureCache()
assert cache.get_session_signature("sid-test") is None
cache.cache_session_signature("sid-test", _SIG_A, 5)
assert cache.get_session_signature("sid-test") == _SIG_A
def test_session_signature_longer_replaces_same_count() -> None:
"""同一 message_count 下，更长的签名替换更短的。"""
cache = ThinkingSignatureCache()
sig_short = "x" * 60
sig_long = "y" * 80
cache.cache_session_signature("sid-1", sig_short, 5)
cache.cache_session_signature("sid-1", sig_long, 5)
assert cache.get_session_signature("sid-1") == sig_long
# 更短的不会替换
cache.cache_session_signature("sid-1", sig_short, 5)
assert cache.get_session_signature("sid-1") == sig_long
def test_session_signature_rewind_detection() -> None:
"""Rewind: message_count 减少时强制更新签名。"""
cache = ThinkingSignatureCache()
cache.cache_session_signature("sid-1", _SIG_A, 10)
assert cache.get_session_signature("sid-1") == _SIG_A
# message_count 3 < 10 → rewind detected, force update
cache.cache_session_signature("sid-1", _SIG_B, 3)
assert cache.get_session_signature("sid-1") == _SIG_B
def test_session_signature_short_ignored() -> None:
"""短签名即使 rewind 也不会被缓存。"""
cache = ThinkingSignatureCache()
cache.cache_session_signature("sid-1", _SIG_A, 5)
cache.cache_session_signature("sid-1", "short", 1)
assert cache.get_session_signature("sid-1") == _SIG_A
def test_session_isolation() -> None:
"""不同 session 之间互相隔离。"""
cache = ThinkingSignatureCache()
cache.cache_session_signature("sid-1", _SIG_A, 1)
assert cache.get_session_signature("sid-1") == _SIG_A
assert cache.get_session_signature("sid-2") is None
# ===== Clear =====
def test_clear_all_layers() -> None:
cache = ThinkingSignatureCache()
cache.cache_tool_signature("tool-1", _SIG_A)
cache.cache_thinking_family(_SIG_A, "claude")
cache.cache_session_signature("sid-1", _SIG_B, 1)
cache.cache("gemini-3-pro", "text", _SIG_C)
cache.clear()
assert cache.get_tool_signature("tool-1") is None
assert cache.get_signature_family(_SIG_A) is None
assert cache.get_session_signature("sid-1") is None
# Legacy layer: 回退到 DUMMY
assert cache.get_or_dummy("gemini-3-pro", "text") == DUMMY_THOUGHT_SIGNATURE

11
tests/services/antigravity/test_thinking_signature_conversion.py
View File

@@ -1,16 +1,18 @@
from __future__ import annotations
from src.api.handlers.base.utils import get_format_converter_registry
from src.services.antigravity.constants import DUMMY_THOUGHT_SIGNATURE
from src.services.antigravity.signature_cache import signature_cache
from src.services.provider.adapters.antigravity.constants import DUMMY_THOUGHT_SIGNATURE
from src.services.provider.adapters.antigravity.signature_cache import signature_cache
def _reset_sig_cache() -> None:
# Module-global cache; tests must isolate state.
signature_cache._cache.clear() # type: ignore[attr-defined]
signature_cache.clear()
def test_antigravity_converts_claude_thinking_block_to_gemini_thought_part_prefers_payload_sig() -> None:
def test_antigravity_converts_claude_thinking_block_to_gemini_thought_part_prefers_payload_sig() -> (
None
):
_reset_sig_cache()
req = {
@@ -115,4 +117,3 @@ def test_antigravity_inserts_dummy_thought_for_last_assistant_when_thinking_enab
assert parts[0]["thought"] is True
assert parts[0]["thoughtSignature"] == DUMMY_THOUGHT_SIGNATURE
assert parts[1]["text"] == "prefill"

22
tests/services/antigravity/test_url_availability.py
View File

@@ -1,12 +1,15 @@
from __future__ import annotations
import src.services.antigravity.url_availability as ua_mod
from src.services.antigravity.constants import (
from typing import Any
import src.services.provider.adapters.antigravity.url_availability as ua_mod
from src.services.provider.adapters.antigravity.constants import (
DAILY_BASE_URL,
PROD_BASE_URL,
SANDBOX_BASE_URL,
URL_UNAVAILABLE_TTL_SECONDS,
)
from src.services.antigravity.url_availability import url_availability
from src.services.provider.adapters.antigravity.url_availability import url_availability
def _reset_state() -> None:
@@ -31,11 +34,13 @@ def test_mark_unavailable_filters_url() -> None:
url_availability.mark_unavailable(DAILY_BASE_URL)
ordered = url_availability.get_ordered_urls(prefer_daily=True)
assert ordered[0] == PROD_BASE_URL
# Sandbox → Prod（Daily 被标记为不可用，应被过滤掉）
assert ordered[0] == SANDBOX_BASE_URL
assert DAILY_BASE_URL not in ordered
assert url_availability.is_available(DAILY_BASE_URL) is False
def test_ttl_recovery(monkeypatch) -> None:
def test_ttl_recovery(monkeypatch: Any) -> None:
_reset_state()
t0 = 1000.0
@@ -50,9 +55,10 @@ def test_ttl_recovery(monkeypatch) -> None:
def test_all_unavailable_fallback_returns_base_order() -> None:
_reset_state()
url_availability.mark_unavailable(PROD_BASE_URL)
url_availability.mark_unavailable(SANDBOX_BASE_URL)
url_availability.mark_unavailable(DAILY_BASE_URL)
url_availability.mark_unavailable(PROD_BASE_URL)
ordered = url_availability.get_ordered_urls(prefer_daily=True)
assert ordered == [DAILY_BASE_URL, PROD_BASE_URL]
# 全部不可用时仍返回基础顺序（允许继续尝试，等 TTL 恢复）
assert ordered == [SANDBOX_BASE_URL, DAILY_BASE_URL, PROD_BASE_URL]

61
tests/services/test_claude_fetch_models_pagination.py Normal file
View File

@@ -0,0 +1,61 @@
from __future__ import annotations
from dataclasses import dataclass
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from src.api.handlers.claude.adapter import ClaudeChatAdapter
@dataclass
class _DummyResp:
status_code: int
payload: object
text: str = ""
def json(self) -> object: # noqa: D401
"""Return mocked JSON payload."""
return self.payload
@pytest.mark.asyncio
async def test_claude_fetch_models_paginates_until_has_more_false() -> None:
page1 = {
"data": [{"id": "m1"}, {"id": "m2"}],
"has_more": True,
"first_id": "m1",
"last_id": "m2",
}
page2 = {
"data": [{"id": "m3"}],
"has_more": False,
"first_id": "m3",
"last_id": "m3",
}
client = SimpleNamespace(
get=AsyncMock(
side_effect=[
_DummyResp(status_code=200, payload=page1),
_DummyResp(status_code=200, payload=page2),
]
)
)
models, err = await ClaudeChatAdapter.fetch_models(
client, # type: ignore[arg-type]
"https://api.anthropic.com",
"k",
None,
)
assert err is None
assert [m.get("id") for m in models] == ["m1", "m2", "m3"]
assert all(m.get("api_format") == "claude:chat" for m in models)
# Second page should pass after_id=last_id from page1.
assert client.get.call_count == 2
_, kwargs2 = client.get.call_args_list[1]
assert kwargs2.get("params", {}).get("after_id") == "m2"

4
tests/services/test_codex_request.py
View File

@@ -1,6 +1,6 @@
from __future__ import annotations
from src.services.provider.codex import (
from src.services.provider.adapters.codex.request_patching import (
maybe_patch_request_for_codex,
patch_openai_cli_request_for_codex,
)
@@ -106,7 +106,7 @@ def test_maybe_patch_request_for_codex_patches_for_codex_openai_cli() -> None:
def test_codex_envelope_extra_headers_includes_sse_accept_and_session() -> None:
from src.services.codex.envelope import codex_oauth_envelope
from src.services.provider.adapters.codex.envelope import codex_oauth_envelope
headers = codex_oauth_envelope.extra_headers() or {}
assert headers.get("Accept") == "text/event-stream"

9
tests/services/test_provider_transport_antigravity.py
View File

@@ -4,7 +4,7 @@ from dataclasses import dataclass
from types import SimpleNamespace
from unittest.mock import patch
from src.services.antigravity.constants import PROD_BASE_URL
from src.services.provider.adapters.antigravity.constants import PROD_BASE_URL
from src.services.provider.transport import build_provider_url, get_antigravity_base_url
@@ -24,11 +24,11 @@ def test_antigravity_uses_v1internal_path_and_sets_contextvar() -> None:
)
with patch(
"src.services.provider.transport.url_availability.get_ordered_urls",
"src.services.provider.adapters.antigravity.plugin.url_availability.get_ordered_urls",
return_value=[PROD_BASE_URL],
):
url = build_provider_url(
endpoint, # type: ignore[arg-type] - test stub
endpoint, # type: ignore[arg-type]
path_params={"model": "gemini-2.0-flash"},
is_stream=True,
)
@@ -46,11 +46,10 @@ def test_gemini_cli_non_antigravity_uses_v1beta_path_and_clears_contextvar() ->
)
url = build_provider_url(
endpoint, # type: ignore[arg-type] - test stub
endpoint, # type: ignore[arg-type]
path_params={"model": "gemini-2.0-flash"},
is_stream=False,
)
assert "/v1beta/models/gemini-2.0-flash:generateContent" in url
assert get_antigravity_base_url() is None

5
tests/services/test_provider_transport_codex.py
View File

@@ -22,7 +22,7 @@ def test_codex_openai_cli_uses_responses_path_without_v1_prefix() -> None:
)
url = build_provider_url(
endpoint, # type: ignore[arg-type] - test stub
endpoint, # type: ignore[arg-type]
path_params={"model": "ignored"},
is_stream=True,
)
@@ -38,10 +38,9 @@ def test_codex_openai_cli_does_not_duplicate_responses_suffix() -> None:
)
url = build_provider_url(
endpoint, # type: ignore[arg-type] - test stub
endpoint, # type: ignore[arg-type]
path_params={"model": "ignored"},
is_stream=False,
)
assert url == "https://chatgpt.com/backend-api/codex/responses"

7
tests/services/test_request_builder_oauth.py
View File

@@ -29,18 +29,19 @@ async def test_get_provider_auth_oauth_returns_decrypted_auth_config() -> None:
provider=None,
)
def _decrypt(v): # noqa: ANN001
def _decrypt(v: str) -> str:
if v == "enc_access":
return "access-token"
if v == "enc_cfg":
return json.dumps(token_meta)
return ""
with patch("src.api.handlers.base.request_builder.crypto_service.decrypt", side_effect=_decrypt):
with patch(
"src.api.handlers.base.request_builder.crypto_service.decrypt", side_effect=_decrypt
):
auth = await get_provider_auth(endpoint, key) # type: ignore[arg-type]
assert auth is not None
assert auth.auth_header == "Authorization"
assert auth.auth_value == "Bearer access-token"
assert auth.decrypted_auth_config == token_meta

69
tests/services/test_upstream_fetcher_antigravity_models.py Normal file
View File

@@ -0,0 +1,69 @@
from __future__ import annotations
from unittest.mock import AsyncMock, patch
import pytest
from src.services.model.upstream_fetcher import UpstreamModelsFetchContext, fetch_models_for_key
@pytest.mark.asyncio
async def test_fetch_models_for_key_antigravity_parses_quota() -> None:
mock_resp = {
"models": {
"claude-sonnet-4": {
"displayName": "Claude Sonnet 4",
"quotaInfo": {"remainingFraction": 0.75, "resetTime": "2024-01-15T12:00:00Z"},
},
"gemini-2.5-pro": {
"displayName": "Gemini 2.5 Pro",
"quotaInfo": {"remainingFraction": 0.0},
},
}
}
ctx = UpstreamModelsFetchContext(
provider_type="antigravity",
api_key_value="tok",
format_to_endpoint={},
proxy_config=None,
auth_config={"project_id": "project-1"},
)
with patch(
"src.services.provider.adapters.antigravity.client.fetch_available_models",
AsyncMock(return_value=mock_resp),
):
models, errors, ok, meta = await fetch_models_for_key(ctx, timeout_seconds=1.0)
assert ok is True
assert errors == []
ids = {m.get("id") for m in models}
assert "claude-sonnet-4" in ids
assert "gemini-2.5-pro" in ids
assert isinstance(meta, dict)
quota = meta["antigravity"]["quota_by_model"]
assert quota["claude-sonnet-4"]["remaining_fraction"] == 0.75
assert quota["claude-sonnet-4"]["used_percent"] == 25.0
assert quota["claude-sonnet-4"]["reset_time"] == "2024-01-15T12:00:00Z"
assert quota["gemini-2.5-pro"]["used_percent"] == 100.0
@pytest.mark.asyncio
async def test_fetch_models_for_key_antigravity_requires_project_id() -> None:
ctx = UpstreamModelsFetchContext(
provider_type="antigravity",
api_key_value="tok",
format_to_endpoint={},
proxy_config=None,
auth_config={},
)
models, errors, ok, meta = await fetch_models_for_key(ctx, timeout_seconds=1.0)
assert ok is False
assert models == []
assert meta is None
assert any("project_id" in e for e in errors)

5
tests/services/test_upstream_stream_policy.py
View File

@@ -18,7 +18,9 @@ class _DummyEndpoint:
def test_get_upstream_stream_policy_defaults_to_auto() -> None:
ep = _DummyEndpoint(api_format="openai:chat", config=None, provider=SimpleNamespace(provider_type="custom"))
ep = _DummyEndpoint(
api_format="openai:chat", config=None, provider=SimpleNamespace(provider_type="custom")
)
assert get_upstream_stream_policy(ep) == UpstreamStreamPolicy.AUTO
@@ -60,4 +62,3 @@ def test_enforce_stream_mode_for_upstream_gemini_drops_stream_field() -> None:
)
assert "stream" not in out
assert out["foo"] == "bar"

299
tests/services/test_usage_queue_events.py
View File

@@ -1,5 +1,6 @@
import asyncio
import json
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@@ -25,10 +26,16 @@ from src.services.usage.telemetry_writer import (
class DummyRedis:
def __init__(self) -> None:
self.calls = []
self.xadd_error = None
self.calls: list[tuple[str, dict[str, str], int | None, bool | None]] = []
self.xadd_error: Exception | None = None
async def xadd(self, key, fields, maxlen=None, approximate=None):
async def xadd(
self,
key: str,
fields: dict[str, str],
maxlen: int | None = None,
approximate: bool | None = None,
) -> str:
if self.xadd_error:
raise self.xadd_error
self.calls.append((key, fields, maxlen, approximate))
@@ -36,7 +43,7 @@ class DummyRedis:
@pytest.mark.asyncio
async def test_usage_event_roundtrip():
async def test_usage_event_roundtrip() -> None:
event = build_usage_event(
event_type=UsageEventType.COMPLETED,
request_id="req-1",
@@ -52,10 +59,10 @@ async def test_usage_event_roundtrip():
@pytest.mark.asyncio
async def test_queue_writer_publishes_event(monkeypatch):
async def test_queue_writer_publishes_event(monkeypatch: Any) -> None:
dummy = DummyRedis()
async def _get_redis_client(require_redis=False):
async def _get_redis_client(require_redis: bool = False) -> Any:
return dummy
monkeypatch.setattr("src.services.usage.telemetry_writer.get_redis_client", _get_redis_client)
@@ -96,7 +103,7 @@ async def test_queue_writer_publishes_event(monkeypatch):
@pytest.mark.asyncio
async def test_usage_event_all_types():
async def test_usage_event_all_types() -> None:
"""测试所有事件类型的序列化/反序列化"""
for event_type in UsageEventType:
event = build_usage_event(
@@ -111,7 +118,7 @@ async def test_usage_event_all_types():
@pytest.mark.asyncio
async def test_usage_event_bytes_payload():
async def test_usage_event_bytes_payload() -> None:
"""测试 bytes 类型 payload 的反序列化"""
event = build_usage_event(
event_type=UsageEventType.COMPLETED,
@@ -120,19 +127,19 @@ async def test_usage_event_bytes_payload():
)
fields = event.to_stream_fields()
# 模拟 Redis 返回 bytes
fields["payload"] = fields["payload"].encode("utf-8")
fields["payload"] = fields["payload"].encode("utf-8") # type: ignore[assignment]
restored = UsageEvent.from_stream_fields(fields)
assert restored.request_id == "req-bytes"
@pytest.mark.asyncio
async def test_usage_event_missing_payload():
async def test_usage_event_missing_payload() -> None:
"""测试缺少 payload 字段时抛出异常"""
with pytest.raises(ValueError, match="Missing payload"):
UsageEvent.from_stream_fields({})
def test_sanitize_payload_nested():
def test_sanitize_payload_nested() -> None:
"""测试 sanitize_payload 处理嵌套结构"""
data = {
"str": "hello",
@@ -151,7 +158,7 @@ def test_sanitize_payload_nested():
assert isinstance(result["custom"], str)
def test_parse_body_json_string():
def test_parse_body_json_string() -> None:
"""测试 _parse_body 正确反序列化 JSON 字符串"""
from src.services.usage.consumer_streams import _parse_body
@@ -162,7 +169,7 @@ def test_parse_body_json_string():
assert result["messages"][0]["content"] == "hello"
def test_parse_body_dict_passthrough():
def test_parse_body_dict_passthrough() -> None:
"""测试 _parse_body 直接返回 dict"""
from src.services.usage.consumer_streams import _parse_body
@@ -171,14 +178,14 @@ def test_parse_body_dict_passthrough():
assert result is data # 应该是同一个对象
def test_parse_body_none():
def test_parse_body_none() -> None:
"""测试 _parse_body 处理 None"""
from src.services.usage.consumer_streams import _parse_body
assert _parse_body(None) is None
def test_parse_body_truncated_string():
def test_parse_body_truncated_string() -> None:
"""测试 _parse_body 处理被截断的 JSON 字符串"""
from src.services.usage.consumer_streams import _parse_body
@@ -192,7 +199,7 @@ def test_parse_body_truncated_string():
@pytest.mark.asyncio
async def test_db_telemetry_writer_filters_kwargs():
async def test_db_telemetry_writer_filters_kwargs() -> None:
"""测试 DbTelemetryWriter 过滤不支持的参数"""
mock_telemetry = MagicMock()
mock_telemetry.record_success = AsyncMock()
@@ -219,7 +226,7 @@ async def test_db_telemetry_writer_filters_kwargs():
@pytest.mark.asyncio
async def test_db_telemetry_writer_all_methods():
async def test_db_telemetry_writer_all_methods() -> None:
"""测试 DbTelemetryWriter 的所有方法"""
mock_telemetry = MagicMock()
mock_telemetry.record_success = AsyncMock()
@@ -238,11 +245,11 @@ async def test_db_telemetry_writer_all_methods():
@pytest.mark.asyncio
async def test_queue_writer_record_failure(monkeypatch):
async def test_queue_writer_record_failure(monkeypatch: Any) -> None:
"""测试 QueueTelemetryWriter.record_failure"""
dummy = DummyRedis()
async def _get_redis_client(require_redis=False):
async def _get_redis_client(require_redis: bool = False) -> Any:
return dummy
monkeypatch.setattr("src.services.usage.telemetry_writer.get_redis_client", _get_redis_client)
@@ -275,11 +282,11 @@ async def test_queue_writer_record_failure(monkeypatch):
@pytest.mark.asyncio
async def test_queue_writer_record_cancelled(monkeypatch):
async def test_queue_writer_record_cancelled(monkeypatch: Any) -> None:
"""测试 QueueTelemetryWriter.record_cancelled"""
dummy = DummyRedis()
async def _get_redis_client(require_redis=False):
async def _get_redis_client(require_redis: bool = False) -> Any:
return dummy
monkeypatch.setattr("src.services.usage.telemetry_writer.get_redis_client", _get_redis_client)
@@ -296,11 +303,11 @@ async def test_queue_writer_record_cancelled(monkeypatch):
@pytest.mark.asyncio
async def test_queue_writer_include_headers_bodies(monkeypatch):
async def test_queue_writer_include_headers_bodies(monkeypatch: Any) -> None:
"""测试 include_headers 和 include_bodies 配置"""
dummy = DummyRedis()
async def _get_redis_client(require_redis=False):
async def _get_redis_client(require_redis: bool = False) -> Any:
return dummy
monkeypatch.setattr("src.services.usage.telemetry_writer.get_redis_client", _get_redis_client)
@@ -332,7 +339,7 @@ async def test_queue_writer_include_headers_bodies(monkeypatch):
@pytest.mark.asyncio
async def test_event_to_record_body_deserialization(monkeypatch):
async def test_event_to_record_body_deserialization(monkeypatch: Any) -> None:
"""测试 _event_to_record 正确反序列化 body 字符串为 dict"""
from src.services.usage.consumer_streams import _event_to_record
@@ -361,11 +368,11 @@ async def test_event_to_record_body_deserialization(monkeypatch):
@pytest.mark.asyncio
async def test_queue_writer_body_truncation(monkeypatch):
async def test_queue_writer_body_truncation(monkeypatch: Any) -> None:
"""测试 body 超长截断"""
dummy = DummyRedis()
async def _get_redis_client(require_redis=False):
async def _get_redis_client(require_redis: bool = False) -> Any:
return dummy
monkeypatch.setattr("src.services.usage.telemetry_writer.get_redis_client", _get_redis_client)
@@ -393,10 +400,10 @@ async def test_queue_writer_body_truncation(monkeypatch):
@pytest.mark.asyncio
async def test_queue_writer_redis_unavailable(monkeypatch):
async def test_queue_writer_redis_unavailable(monkeypatch: Any) -> None:
"""测试 Redis 不可用时抛出异常"""
async def _get_redis_client(require_redis=False):
async def _get_redis_client(require_redis: bool = False) -> Any:
return None
monkeypatch.setattr("src.services.usage.telemetry_writer.get_redis_client", _get_redis_client)
@@ -411,12 +418,12 @@ async def test_queue_writer_redis_unavailable(monkeypatch):
@pytest.mark.asyncio
async def test_queue_writer_xadd_error(monkeypatch):
async def test_queue_writer_xadd_error(monkeypatch: Any) -> None:
"""测试 XADD 失败时抛出异常"""
dummy = DummyRedis()
dummy.xadd_error = Exception("Redis connection lost")
async def _get_redis_client(require_redis=False):
async def _get_redis_client(require_redis: bool = False) -> Any:
return dummy
monkeypatch.setattr("src.services.usage.telemetry_writer.get_redis_client", _get_redis_client)
@@ -438,13 +445,13 @@ class MockRedisPipeline:
def __init__(self, parent: "MockRedisForConsumer") -> None:
self._parent = parent
self._commands = []
self._commands: list[tuple[str, ...]] = []
def xack(self, key, group, message_id):
def xack(self, key: str, group: str, message_id: str) -> Any:
self._commands.append(("xack", key, group, message_id))
return self
async def execute(self):
async def execute(self) -> Any:
results = []
for cmd in self._commands:
if cmd[0] == "xack":
@@ -458,54 +465,69 @@ class MockRedisForConsumer:
"""模拟 Redis 客户端用于消费者测试"""
def __init__(self) -> None:
self.xgroup_create_calls = []
self.xreadgroup_results = []
self.xautoclaim_results = []
self.xack_calls = []
self.xadd_calls = []
self.xpending_range_results = []
self.xlen_result = 0
self.xpending_result = {"pending": 0}
self.xgroup_create_error = None
self.xgroup_create_calls: list[tuple[str, str, str, bool]] = []
self.xreadgroup_results: list[Any] = []
self.xautoclaim_results: list[Any] = []
self.xack_calls: list[tuple[str, str, str]] = []
self.xadd_calls: list[tuple[str, dict[str, str], int | None, bool | None]] = []
self.xpending_range_results: list[Any] = []
self.xlen_result: int = 0
self.xpending_result: dict[str, int] | tuple[int, str, str, list[Any]] = {"pending": 0}
self.xgroup_create_error: Exception | None = None
async def xgroup_create(self, key, group, id, mkstream=False):
async def xgroup_create(self, key: str, group: str, id: str, mkstream: bool = False) -> Any:
self.xgroup_create_calls.append((key, group, id, mkstream))
if self.xgroup_create_error:
raise self.xgroup_create_error
async def xreadgroup(self, groupname, consumername, streams, count, block):
async def xreadgroup(
self,
groupname: str,
consumername: str,
streams: dict[str, str],
count: int,
block: int,
) -> Any:
if self.xreadgroup_results:
return self.xreadgroup_results.pop(0)
return None
async def xautoclaim(self, key, group, consumer, min_idle_time, start_id, count):
async def xautoclaim(
self, key: str, group: str, consumer: str, min_idle_time: int, start_id: str, count: int
) -> Any:
if self.xautoclaim_results:
return self.xautoclaim_results.pop(0)
return None
async def xack(self, key, group, message_id):
async def xack(self, key: str, group: str, message_id: str) -> Any:
self.xack_calls.append((key, group, message_id))
async def xadd(self, key, fields, maxlen=None, approximate=None):
async def xadd(
self,
key: str,
fields: dict[str, str],
maxlen: int | None = None,
approximate: bool | None = None,
) -> Any:
self.xadd_calls.append((key, fields, maxlen, approximate))
return "dlq-1-0"
async def xpending_range(self, key, group, min, max, count):
async def xpending_range(self, key: str, group: str, min: str, max: str, count: int) -> Any:
if self.xpending_range_results:
return self.xpending_range_results.pop(0)
return []
async def xlen(self, key):
async def xlen(self, key: str) -> Any:
return self.xlen_result
async def xpending(self, key, group):
async def xpending(self, key: str, group: str) -> Any:
return self.xpending_result
def pipeline(self):
def pipeline(self) -> Any:
return MockRedisPipeline(self)
def test_consumer_name():
def test_consumer_name() -> None:
"""测试消费者名称生成"""
name = _consumer_name()
assert ":" in name
@@ -516,11 +538,11 @@ def test_consumer_name():
@pytest.mark.asyncio
async def test_ensure_stream_group_creates_group(monkeypatch):
async def test_ensure_stream_group_creates_group(monkeypatch: Any) -> None:
"""测试创建消费者组"""
mock_redis = MockRedisForConsumer()
async def _get_redis_client(require_redis=False):
async def _get_redis_client(require_redis: bool = False) -> Any:
return mock_redis
monkeypatch.setattr("src.services.usage.consumer_streams.get_redis_client", _get_redis_client)
@@ -535,12 +557,12 @@ async def test_ensure_stream_group_creates_group(monkeypatch):
@pytest.mark.asyncio
async def test_ensure_stream_group_handles_busygroup(monkeypatch):
async def test_ensure_stream_group_handles_busygroup(monkeypatch: Any) -> None:
"""测试消费者组已存在时不抛异常"""
mock_redis = MockRedisForConsumer()
mock_redis.xgroup_create_error = ResponseError("BUSYGROUP Consumer Group name already exists")
async def _get_redis_client(require_redis=False):
async def _get_redis_client(require_redis: bool = False) -> Any:
return mock_redis
monkeypatch.setattr("src.services.usage.consumer_streams.get_redis_client", _get_redis_client)
@@ -550,12 +572,12 @@ async def test_ensure_stream_group_handles_busygroup(monkeypatch):
@pytest.mark.asyncio
async def test_ensure_stream_group_raises_other_errors(monkeypatch):
async def test_ensure_stream_group_raises_other_errors(monkeypatch: Any) -> None:
"""测试其他 Redis 错误时抛出异常"""
mock_redis = MockRedisForConsumer()
mock_redis.xgroup_create_error = ResponseError("SOME OTHER ERROR")
async def _get_redis_client(require_redis=False):
async def _get_redis_client(require_redis: bool = False) -> Any:
return mock_redis
monkeypatch.setattr("src.services.usage.consumer_streams.get_redis_client", _get_redis_client)
@@ -565,10 +587,10 @@ async def test_ensure_stream_group_raises_other_errors(monkeypatch):
@pytest.mark.asyncio
async def test_ensure_stream_group_no_redis(monkeypatch):
async def test_ensure_stream_group_no_redis(monkeypatch: Any) -> None:
"""测试 Redis 不可用时直接返回"""
async def _get_redis_client(require_redis=False):
async def _get_redis_client(require_redis: bool = False) -> Any:
return None
monkeypatch.setattr("src.services.usage.consumer_streams.get_redis_client", _get_redis_client)
@@ -578,13 +600,13 @@ async def test_ensure_stream_group_no_redis(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_start_stop():
async def test_consumer_start_stop() -> None:
"""测试消费者启动和停止"""
consumer = UsageQueueConsumer()
assert not consumer._running
# Mock _run 避免真正执行
consumer._run = AsyncMock()
consumer._run = AsyncMock() # type: ignore[method-assign]
await consumer.start()
assert consumer._running
@@ -603,7 +625,7 @@ async def test_consumer_start_stop():
@pytest.mark.asyncio
async def test_consumer_process_messages_success(monkeypatch):
async def test_consumer_process_messages_success(monkeypatch: Any) -> None:
"""测试成功处理消息"""
mock_redis = MockRedisForConsumer()
@@ -626,7 +648,7 @@ async def test_consumer_process_messages_success(monkeypatch):
consumer = UsageQueueConsumer()
# Mock 批量处理方法
consumer._process_record_batch = AsyncMock()
consumer._process_record_batch = AsyncMock() # type: ignore[method-assign]
await consumer._process_messages(mock_redis, messages)
@@ -638,7 +660,7 @@ async def test_consumer_process_messages_success(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_process_messages_error_retry(monkeypatch):
async def test_consumer_process_messages_error_retry(monkeypatch: Any) -> None:
"""测试处理消息失败时 STREAMING 事件的重试行为"""
mock_redis = MockRedisForConsumer()
# 返回重试次数小于 max_retries
@@ -660,7 +682,7 @@ async def test_consumer_process_messages_error_retry(monkeypatch):
# 在设置 config 后创建 consumer，以便缓存正确的配置值
consumer = UsageQueueConsumer()
# 模拟 STREAMING 事件处理失败
consumer._apply_streaming_event = AsyncMock(side_effect=ValueError("Processing error"))
consumer._apply_streaming_event = AsyncMock(side_effect=ValueError("Processing error")) # type: ignore[method-assign]
await consumer._process_messages(mock_redis, messages)
@@ -673,7 +695,7 @@ async def test_consumer_process_messages_error_retry(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_process_messages_move_to_dlq(monkeypatch):
async def test_consumer_process_messages_move_to_dlq(monkeypatch: Any) -> None:
"""测试消息超过最大重试次数后移入 DLQ"""
mock_redis = MockRedisForConsumer()
# 返回重试次数 >= max_retries
@@ -712,7 +734,7 @@ async def test_consumer_process_messages_move_to_dlq(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_get_delivery_count_dict_format():
async def test_consumer_get_delivery_count_dict_format() -> None:
"""测试获取消息投递次数（dict 格式）"""
mock_redis = MockRedisForConsumer()
mock_redis.xpending_range_results = [[{"times_delivered": 3}]]
@@ -723,7 +745,7 @@ async def test_consumer_get_delivery_count_dict_format():
@pytest.mark.asyncio
async def test_consumer_get_delivery_count_tuple_format():
async def test_consumer_get_delivery_count_tuple_format() -> None:
"""测试获取消息投递次数（tuple 格式）"""
mock_redis = MockRedisForConsumer()
# (message_id, consumer, idle_time, times_delivered)
@@ -735,7 +757,7 @@ async def test_consumer_get_delivery_count_tuple_format():
@pytest.mark.asyncio
async def test_consumer_get_delivery_count_empty():
async def test_consumer_get_delivery_count_empty() -> None:
"""测试消息不存在时返回 0"""
mock_redis = MockRedisForConsumer()
mock_redis.xpending_range_results = [[]]
@@ -746,7 +768,7 @@ async def test_consumer_get_delivery_count_empty():
@pytest.mark.asyncio
async def test_consumer_read_new_messages(monkeypatch):
async def test_consumer_read_new_messages(monkeypatch: Any) -> None:
"""测试读取新消息"""
mock_redis = MockRedisForConsumer()
@@ -761,7 +783,7 @@ async def test_consumer_read_new_messages(monkeypatch):
]
consumer = UsageQueueConsumer()
consumer._process_messages = AsyncMock()
consumer._process_messages = AsyncMock() # type: ignore[method-assign]
await consumer._read_new(mock_redis)
@@ -772,14 +794,14 @@ async def test_consumer_read_new_messages(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_maybe_claim_pending_respects_interval(monkeypatch):
async def test_consumer_maybe_claim_pending_respects_interval(monkeypatch: Any) -> None:
"""测试 claim 间隔限制"""
mock_redis = MockRedisForConsumer()
consumer = UsageQueueConsumer()
consumer._last_claim = 999999999999.0 # 未来时间
consumer._process_messages = AsyncMock()
consumer._process_messages = AsyncMock() # type: ignore[method-assign]
await consumer._maybe_claim_pending(mock_redis)
@@ -788,14 +810,14 @@ async def test_consumer_maybe_claim_pending_respects_interval(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_maybe_claim_pending_xautoclaim_error(monkeypatch):
async def test_consumer_maybe_claim_pending_xautoclaim_error(monkeypatch: Any) -> None:
"""测试 XAUTOCLAIM 失败时优雅处理"""
mock_redis = MockRedisForConsumer()
async def failing_xautoclaim(*args, **kwargs):
async def failing_xautoclaim(*args: Any, **kwargs: Any) -> Any:
raise ResponseError("XAUTOCLAIM error")
mock_redis.xautoclaim = failing_xautoclaim
mock_redis.xautoclaim = failing_xautoclaim # type: ignore[method-assign]
consumer = UsageQueueConsumer()
consumer._last_claim = 0 # 确保会尝试 claim
@@ -805,7 +827,7 @@ async def test_consumer_maybe_claim_pending_xautoclaim_error(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_log_metrics_dict_pending():
async def test_consumer_log_metrics_dict_pending() -> None:
"""测试 metrics 日志（dict 格式 pending）"""
mock_redis = MockRedisForConsumer()
mock_redis.xlen_result = 100
@@ -825,7 +847,7 @@ async def test_consumer_log_metrics_dict_pending():
@pytest.mark.asyncio
async def test_consumer_log_metrics_tuple_pending():
async def test_consumer_log_metrics_tuple_pending() -> None:
"""测试 metrics 日志（tuple 格式 pending）"""
mock_redis = MockRedisForConsumer()
mock_redis.xlen_result = 50
@@ -845,11 +867,11 @@ async def test_consumer_log_metrics_tuple_pending():
@pytest.mark.asyncio
async def test_consumer_apply_event_streaming(monkeypatch):
async def test_consumer_apply_event_streaming(monkeypatch: Any) -> None:
"""测试处理 STREAMING 事件"""
mock_db = MagicMock()
def mock_create_session():
def mock_create_session() -> Any:
return mock_db
monkeypatch.setattr("src.services.usage.consumer_streams.create_session", mock_create_session)
@@ -884,12 +906,12 @@ async def test_consumer_apply_event_streaming(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_apply_event_completed(monkeypatch):
async def test_consumer_apply_event_completed(monkeypatch: Any) -> None:
"""测试处理 COMPLETED 事件"""
mock_db = MagicMock()
mock_db.query.return_value.filter.return_value.first.return_value = None
def mock_create_session():
def mock_create_session() -> Any:
return mock_db
monkeypatch.setattr("src.services.usage.consumer_streams.create_session", mock_create_session)
@@ -925,12 +947,12 @@ async def test_consumer_apply_event_completed(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_apply_event_failed(monkeypatch):
async def test_consumer_apply_event_failed(monkeypatch: Any) -> None:
"""测试处理 FAILED 事件"""
mock_db = MagicMock()
mock_db.query.return_value.filter.return_value.first.return_value = None
def mock_create_session():
def mock_create_session() -> Any:
return mock_db
monkeypatch.setattr("src.services.usage.consumer_streams.create_session", mock_create_session)
@@ -962,12 +984,12 @@ async def test_consumer_apply_event_failed(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_apply_event_cancelled(monkeypatch):
async def test_consumer_apply_event_cancelled(monkeypatch: Any) -> None:
"""测试处理 CANCELLED 事件"""
mock_db = MagicMock()
mock_db.query.return_value.filter.return_value.first.return_value = None
def mock_create_session():
def mock_create_session() -> Any:
return mock_db
monkeypatch.setattr("src.services.usage.consumer_streams.create_session", mock_create_session)
@@ -997,7 +1019,7 @@ async def test_consumer_apply_event_cancelled(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_process_messages_batch(monkeypatch):
async def test_consumer_process_messages_batch(monkeypatch: Any) -> None:
"""测试批量处理消息"""
mock_redis = MockRedisForConsumer()
@@ -1023,7 +1045,7 @@ async def test_consumer_process_messages_batch(monkeypatch):
consumer = UsageQueueConsumer()
# Mock 批量处理
consumer._process_record_batch = AsyncMock()
consumer._process_record_batch = AsyncMock() # type: ignore[method-assign]
await consumer._process_messages(mock_redis, messages)
@@ -1032,7 +1054,7 @@ async def test_consumer_process_messages_batch(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_process_messages_separates_streaming(monkeypatch):
async def test_consumer_process_messages_separates_streaming(monkeypatch: Any) -> None:
"""测试消息分类：STREAMING 和其他事件分开处理"""
mock_redis = MockRedisForConsumer()
@@ -1054,8 +1076,8 @@ async def test_consumer_process_messages_separates_streaming(monkeypatch):
]
consumer = UsageQueueConsumer()
consumer._apply_streaming_event = AsyncMock()
consumer._process_record_batch = AsyncMock()
consumer._apply_streaming_event = AsyncMock() # type: ignore[method-assign]
consumer._process_record_batch = AsyncMock() # type: ignore[method-assign]
await consumer._process_messages(mock_redis, messages)
@@ -1072,12 +1094,12 @@ async def test_consumer_process_messages_separates_streaming(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_process_record_batch_success(monkeypatch):
async def test_consumer_process_record_batch_success(monkeypatch: Any) -> None:
"""测试批量记录处理成功"""
mock_redis = MockRedisForConsumer()
mock_db = MagicMock()
def mock_create_session():
def mock_create_session() -> Any:
return mock_db
monkeypatch.setattr("src.services.usage.consumer_streams.create_session", mock_create_session)
@@ -1112,13 +1134,13 @@ async def test_consumer_process_record_batch_success(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_process_record_batch_fallback(monkeypatch):
async def test_consumer_process_record_batch_fallback(monkeypatch: Any) -> None:
"""测试批量处理失败时回退到逐条处理"""
mock_redis = MockRedisForConsumer()
mock_db = MagicMock()
mock_db.query.return_value.filter.return_value.first.return_value = None
def mock_create_session():
def mock_create_session() -> Any:
return mock_db
monkeypatch.setattr("src.services.usage.consumer_streams.create_session", mock_create_session)
@@ -1157,11 +1179,11 @@ async def test_consumer_process_record_batch_fallback(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_apply_streaming_event(monkeypatch):
async def test_consumer_apply_streaming_event(monkeypatch: Any) -> None:
"""测试 STREAMING 事件处理"""
mock_db = MagicMock()
def mock_create_session():
def mock_create_session() -> Any:
return mock_db
monkeypatch.setattr("src.services.usage.consumer_streams.create_session", mock_create_session)
@@ -1195,12 +1217,12 @@ async def test_consumer_apply_streaming_event(monkeypatch):
@pytest.mark.asyncio
async def test_consumer_apply_record_event(monkeypatch):
async def test_consumer_apply_record_event(monkeypatch: Any) -> None:
"""测试记录事件单独处理"""
mock_db = MagicMock()
mock_db.query.return_value.filter.return_value.first.return_value = None
def mock_create_session():
def mock_create_session() -> Any:
return mock_db
monkeypatch.setattr("src.services.usage.consumer_streams.create_session", mock_create_session)
@@ -1226,3 +1248,80 @@ async def test_consumer_apply_record_event(monkeypatch):
call_kwargs = mock_record_usage.call_args.kwargs
assert call_kwargs["request_id"] == "req-record"
assert call_kwargs["input_tokens"] == 100
@pytest.mark.asyncio
async def test_record_usage_batch_updates_when_status_completed_billing_pending(
monkeypatch: Any,
) -> None:
"""回归测试：
usage-queue 模式下，handler 可能会先把 Usage.status 直接更新为 completed（为减少 UI 延迟），
但 billing_status 仍为 pending。此时 completed 事件仍应更新详情字段（如 response_headers/body），
并将 billing_status 结算为 settled。
"""
from src.models.database import Usage
from src.services.usage.service import UsageService
class DummyQuery:
def __init__(self, all_result: list[Any]) -> None:
self._all_result = all_result
def filter(self, *args: Any, **kwargs: Any) -> "DummyQuery":
return self
def all(self) -> list[Any]:
return self._all_result
existing = Usage(
request_id="req-usage-batch-1",
provider_name="pending",
model="gemini-3-pro-image-preview",
status="completed",
billing_status="pending",
)
assert existing.status == "completed"
assert existing.billing_status == "pending"
assert existing.finalized_at is None
db = MagicMock()
db.query.side_effect = lambda model: (
DummyQuery([existing]) if model is Usage else DummyQuery([])
)
usage_params = {
"status": "completed",
"response_headers": {"content-type": "text/event-stream"},
"response_body": {"chunks": [{"foo": "bar"}], "metadata": {"stream": True}},
}
monkeypatch.setattr(
UsageService,
"_prepare_usage_records_batch",
AsyncMock(return_value=[(usage_params, 0.0, None)]),
)
def _fake_update(existing_usage: Any, params: dict[str, Any], _target_model: Any) -> None:
existing_usage.status = params.get("status", existing_usage.status)
existing_usage.response_headers = params.get("response_headers")
existing_usage.response_body = params.get("response_body")
monkeypatch.setattr(UsageService, "_update_existing_usage", _fake_update)
result = await UsageService.record_usage_batch(
db,
[
{
"request_id": "req-usage-batch-1",
"provider": "Antigravity反代",
"model": "gemini-3-pro-image-preview",
"status": "completed",
}
],
)
assert result and result[0] is existing
assert existing.response_headers == usage_params["response_headers"]
assert existing.response_body == usage_params["response_body"]
assert existing.billing_status == "settled"
assert existing.finalized_at is not None