mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
feat(oauth): 账号封禁前置 OAuth 验证、抽取 provider_context、完善账号状态分类
- 新增 verify_oauth_before_account_block:在标记账号封禁前先尝试刷新 token, 区分 OAuth 过期与真正的账号级封禁,避免误标 - 抽取 provider_context.py 统一解析 provider_type,解决 ORM detached 访问问题 - account_state 新增 workspace_deactivated 分类和 auto-removable 状态集合, 补充中文验证关键词匹配 - OAuth refresh 成功后仅清除可恢复的 token 错误,不再自动清除账号级 block - deploy.sh 依赖指纹改用纯 shell 实现,移除对 Python tomllib 的依赖 - 前端 Pool 管理页面新增筛选和批量操作优化 - 补充对应测试用例
This commit is contained in:
108
tests/unit/test_pool_management_filters.py
Normal file
108
tests/unit/test_pool_management_filters.py
Normal file
@@ -0,0 +1,108 @@
|
||||
from src.api.admin.pool.routes import _detail_is_oauth_invalid, _filter_pool_key_details
|
||||
from src.api.admin.pool.schemas import PoolKeyDetail
|
||||
|
||||
|
||||
def _detail(
|
||||
key_id: str,
|
||||
*,
|
||||
is_active: bool = True,
|
||||
scheduling_status: str = "available",
|
||||
account_status_blocked: bool = False,
|
||||
account_status_code: str | None = None,
|
||||
account_status_label: str | None = None,
|
||||
account_status_reason: str | None = None,
|
||||
auth_type: str = "api_key",
|
||||
oauth_invalid_at: int | None = None,
|
||||
oauth_invalid_reason: str | None = None,
|
||||
oauth_expires_at: int | None = None,
|
||||
cooldown_reason: str | None = None,
|
||||
circuit_breaker_open: bool = False,
|
||||
cost_limit: int | None = None,
|
||||
cost_window_usage: int = 0,
|
||||
) -> PoolKeyDetail:
|
||||
return PoolKeyDetail(
|
||||
key_id=key_id,
|
||||
key_name=key_id,
|
||||
is_active=is_active,
|
||||
auth_type=auth_type,
|
||||
oauth_invalid_at=oauth_invalid_at,
|
||||
oauth_invalid_reason=oauth_invalid_reason,
|
||||
oauth_expires_at=oauth_expires_at,
|
||||
scheduling_status=scheduling_status,
|
||||
scheduling_reason=scheduling_status or "available",
|
||||
scheduling_label=scheduling_status or "available",
|
||||
account_status_blocked=account_status_blocked,
|
||||
account_status_code=account_status_code,
|
||||
account_status_label=account_status_label,
|
||||
account_status_reason=account_status_reason,
|
||||
cooldown_reason=cooldown_reason,
|
||||
circuit_breaker_open=circuit_breaker_open,
|
||||
cost_limit=cost_limit,
|
||||
cost_window_usage=cost_window_usage,
|
||||
)
|
||||
|
||||
|
||||
def test_filter_pool_key_details_require_schedulable_keeps_available_and_degraded() -> None:
|
||||
details = [
|
||||
_detail("available", scheduling_status="available"),
|
||||
_detail("degraded", scheduling_status="degraded"),
|
||||
_detail("blocked", scheduling_status="blocked", account_status_blocked=True),
|
||||
]
|
||||
|
||||
filtered = _filter_pool_key_details(details, require_schedulable=True)
|
||||
|
||||
assert [item.key_id for item in filtered] == ["available", "degraded"]
|
||||
|
||||
|
||||
def test_filter_pool_key_details_require_schedulable_uses_fallback_when_status_missing() -> None:
|
||||
details = [
|
||||
_detail("manual-disabled", scheduling_status="", is_active=False),
|
||||
_detail("cooldown", scheduling_status="", cooldown_reason="rate_limited_429"),
|
||||
_detail("usable", scheduling_status="", is_active=True),
|
||||
]
|
||||
|
||||
filtered = _filter_pool_key_details(details, require_schedulable=True)
|
||||
|
||||
assert [item.key_id for item in filtered] == ["usable"]
|
||||
|
||||
|
||||
def test_detail_is_oauth_invalid_excludes_account_disabled_state() -> None:
|
||||
detail = _detail(
|
||||
"disabled-account",
|
||||
auth_type="oauth",
|
||||
oauth_invalid_at=1,
|
||||
oauth_invalid_reason="[ACCOUNT_BLOCK] account has been deactivated",
|
||||
account_status_blocked=True,
|
||||
account_status_code="account_disabled",
|
||||
account_status_label="账号停用",
|
||||
)
|
||||
|
||||
assert _detail_is_oauth_invalid(detail) is False
|
||||
|
||||
|
||||
def test_detail_is_oauth_invalid_accepts_token_expired_state() -> None:
|
||||
detail = _detail(
|
||||
"expired-token",
|
||||
auth_type="oauth",
|
||||
oauth_invalid_at=1,
|
||||
oauth_invalid_reason="[OAUTH_EXPIRED] token invalidated",
|
||||
account_status_blocked=True,
|
||||
account_status_code="oauth_expired",
|
||||
account_status_label="Token 失效",
|
||||
)
|
||||
|
||||
assert _detail_is_oauth_invalid(detail) is True
|
||||
|
||||
|
||||
def test_detail_is_oauth_invalid_accepts_refresh_failed_state() -> None:
|
||||
detail = _detail(
|
||||
"refresh-failed",
|
||||
auth_type="oauth",
|
||||
oauth_invalid_at=1,
|
||||
oauth_invalid_reason="[REFRESH_FAILED] refresh_token_reused",
|
||||
account_status_blocked=False,
|
||||
account_status_code="oauth_refresh_failed",
|
||||
account_status_label="续期失败",
|
||||
)
|
||||
|
||||
assert _detail_is_oauth_invalid(detail) is True
|
||||
@@ -1,7 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -84,6 +86,30 @@ def _make_oauth_key(*, key_id: str, name: str, auth_config: dict[str, object]) -
|
||||
)
|
||||
|
||||
|
||||
class _SingleKeyQuery:
|
||||
def __init__(self, key: SimpleNamespace | None) -> None:
|
||||
self._key = key
|
||||
|
||||
def filter(self, *_args: object, **_kwargs: object) -> "_SingleKeyQuery":
|
||||
return self
|
||||
|
||||
def first(self) -> SimpleNamespace | None:
|
||||
return self._key
|
||||
|
||||
|
||||
class _SingleKeyDB:
|
||||
def __init__(self, key: SimpleNamespace | None) -> None:
|
||||
self._key = key
|
||||
|
||||
def query(self, _model: object) -> _SingleKeyQuery:
|
||||
return _SingleKeyQuery(self._key)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _fake_db_context(db: _SingleKeyDB):
|
||||
yield db
|
||||
|
||||
|
||||
def test_check_duplicate_oauth_account_codex_allows_same_user_different_account_id(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
@@ -151,3 +177,87 @@ def test_check_duplicate_oauth_account_codex_rejects_same_account_user_identity(
|
||||
"plan_type": "team",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_mark_refresh_failed_sync_preserves_existing_account_block(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
key = SimpleNamespace(
|
||||
id="key-1",
|
||||
oauth_invalid_at="old-invalid-at",
|
||||
oauth_invalid_reason="[ACCOUNT_BLOCK] account has been deactivated",
|
||||
)
|
||||
db = _SingleKeyDB(key)
|
||||
monkeypatch.setattr(module, "get_db_context", lambda: _fake_db_context(db))
|
||||
|
||||
module._mark_refresh_failed_sync(
|
||||
"key-1",
|
||||
"[REFRESH_FAILED] Token 续期失败 (400): refresh_token_reused",
|
||||
)
|
||||
|
||||
assert key.oauth_invalid_at == "old-invalid-at"
|
||||
assert key.oauth_invalid_reason == "[ACCOUNT_BLOCK] account has been deactivated"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_account_state_after_oauth_update_refreshes_supported_provider(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
fake_db = SimpleNamespace(close=MagicMock())
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
async def _fake_refresh_provider_quota_for_provider(**kwargs: object) -> dict[str, object]:
|
||||
captured.update(kwargs)
|
||||
return {"success": 1}
|
||||
|
||||
monkeypatch.setattr(module, "create_session", lambda: fake_db)
|
||||
|
||||
from src.services.provider_keys import key_quota_service as quota_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
quota_module,
|
||||
"refresh_provider_quota_for_provider",
|
||||
_fake_refresh_provider_quota_for_provider,
|
||||
)
|
||||
|
||||
attempted, error = await module._refresh_account_state_after_oauth_update(
|
||||
provider_id="provider-1",
|
||||
provider_type="codex",
|
||||
key_ids=["key-1"],
|
||||
)
|
||||
|
||||
assert attempted is True
|
||||
assert error is None
|
||||
assert captured["provider_id"] == "provider-1"
|
||||
assert captured["key_ids"] == ["key-1"]
|
||||
fake_db.close.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_account_state_after_oauth_update_returns_error_when_refresh_fails(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
fake_db = SimpleNamespace(close=MagicMock())
|
||||
|
||||
async def _fake_refresh_provider_quota_for_provider(**_kwargs: object) -> dict[str, object]:
|
||||
raise RuntimeError("quota refresh failed")
|
||||
|
||||
monkeypatch.setattr(module, "create_session", lambda: fake_db)
|
||||
|
||||
from src.services.provider_keys import key_quota_service as quota_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
quota_module,
|
||||
"refresh_provider_quota_for_provider",
|
||||
_fake_refresh_provider_quota_for_provider,
|
||||
)
|
||||
|
||||
attempted, error = await module._refresh_account_state_after_oauth_update(
|
||||
provider_id="provider-1",
|
||||
provider_type="codex",
|
||||
key_ids=["key-1"],
|
||||
)
|
||||
|
||||
assert attempted is True
|
||||
assert "quota refresh failed" in error
|
||||
fake_db.close.assert_called_once()
|
||||
|
||||
@@ -1,18 +1,23 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from src.api.admin.provider_query import TestModelFailoverRequest as FailoverRequestModel
|
||||
from src.api.admin import provider_query as provider_query_module
|
||||
from src.api.admin.provider_query import (
|
||||
DEFAULT_MODEL_TEST_MESSAGE,
|
||||
)
|
||||
from src.api.admin.provider_query import TestModelFailoverRequest as FailoverRequestModel
|
||||
from src.api.admin.provider_query import (
|
||||
_build_direct_test_candidates,
|
||||
_build_test_attempts_from_candidate_keys,
|
||||
_filter_test_candidates_by_endpoint,
|
||||
_flatten_test_candidates_for_concurrency,
|
||||
_maybe_mark_test_oauth_key_invalid,
|
||||
_require_test_endpoint_base_url,
|
||||
_resolve_test_message,
|
||||
_resolve_test_effective_model,
|
||||
_resolve_test_message,
|
||||
)
|
||||
from src.services.scheduling.schemas import PoolCandidate
|
||||
|
||||
@@ -176,3 +181,68 @@ def test_require_test_endpoint_base_url_trims_whitespace() -> None:
|
||||
)
|
||||
|
||||
assert _require_test_endpoint_base_url(endpoint) == "https://api.anthropic.com"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_maybe_mark_test_oauth_key_invalid_skips_account_block_when_oauth_check_fails(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
key = SimpleNamespace(id="key-1", oauth_invalid_at=None, oauth_invalid_reason=None)
|
||||
endpoint = SimpleNamespace(api_format="openai:chat")
|
||||
db = MagicMock()
|
||||
|
||||
async def _fake_verify(**_: object) -> bool:
|
||||
key.oauth_invalid_reason = "[OAUTH_EXPIRED] refresh token expired"
|
||||
return False
|
||||
|
||||
monkeypatch.setattr(provider_query_module, "verify_oauth_before_account_block", _fake_verify)
|
||||
|
||||
await _maybe_mark_test_oauth_key_invalid(
|
||||
db=db,
|
||||
endpoint=endpoint,
|
||||
key=key,
|
||||
auth_type="oauth",
|
||||
error_payload={
|
||||
"error": {
|
||||
"code": 403,
|
||||
"message": "Please verify your account",
|
||||
"status": "PERMISSION_DENIED",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
assert key.oauth_invalid_at is None
|
||||
assert key.oauth_invalid_reason == "[OAUTH_EXPIRED] refresh token expired"
|
||||
db.commit.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_maybe_mark_test_oauth_key_invalid_marks_account_block_after_oauth_check(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
key = SimpleNamespace(id="key-2", oauth_invalid_at=None, oauth_invalid_reason=None)
|
||||
endpoint = SimpleNamespace(api_format="openai:chat")
|
||||
db = MagicMock()
|
||||
|
||||
async def _fake_verify(**_: object) -> bool:
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(provider_query_module, "verify_oauth_before_account_block", _fake_verify)
|
||||
|
||||
await _maybe_mark_test_oauth_key_invalid(
|
||||
db=db,
|
||||
endpoint=endpoint,
|
||||
key=key,
|
||||
auth_type="oauth",
|
||||
error_payload={
|
||||
"error": {
|
||||
"code": 403,
|
||||
"message": "verify your account",
|
||||
"status": "PERMISSION_DENIED",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
assert key.oauth_invalid_at is not None
|
||||
assert str(key.oauth_invalid_reason).startswith("[ACCOUNT_BLOCK] ")
|
||||
db.commit.assert_called_once()
|
||||
|
||||
Reference in New Issue
Block a user