mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
fix(vertex): SA 认证注入代理配置,细化 token 获取异常处理
- _auth_service_account 接收 endpoint 参数,通过 _get_proxy_config 解析代理 - vertex_auth 区分 TimeoutException/RequestError/通用异常,提供可读错误信息 - 新增测试覆盖代理传递和超时场景
This commit is contained in:
@@ -174,8 +174,16 @@ class VertexAuthService:
|
|||||||
raise VertexAuthError(
|
raise VertexAuthError(
|
||||||
f"Failed to get access token: HTTP {e.response.status_code}: {error_body}"
|
f"Failed to get access token: HTTP {e.response.status_code}: {error_body}"
|
||||||
)
|
)
|
||||||
|
except httpx.TimeoutException:
|
||||||
|
raw = client_kwargs.get("timeout")
|
||||||
|
suffix = f" after {raw}s" if isinstance(raw, (int, float)) else ""
|
||||||
|
raise VertexAuthError(f"Failed to get access token: request timed out{suffix}")
|
||||||
|
except httpx.RequestError as e:
|
||||||
|
detail = str(e).strip() or type(e).__name__
|
||||||
|
raise VertexAuthError(f"Failed to get access token: {detail}")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise VertexAuthError(f"Failed to get access token: {e}")
|
detail = str(e).strip() or type(e).__name__
|
||||||
|
raise VertexAuthError(f"Failed to get access token: {detail}")
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def clear_cache(cls, client_email: str | None = None) -> None:
|
def clear_cache(cls, client_email: str | None = None) -> None:
|
||||||
|
|||||||
@@ -12,7 +12,7 @@ from typing import Any
|
|||||||
from src.core.provider_auth_types import ProviderAuthInfo
|
from src.core.provider_auth_types import ProviderAuthInfo
|
||||||
|
|
||||||
|
|
||||||
async def _auth_service_account(key: Any) -> ProviderAuthInfo:
|
async def _auth_service_account(key: Any, endpoint: Any | None = None) -> ProviderAuthInfo:
|
||||||
"""Service Account 认证:SA JSON → JWT → Access Token。"""
|
"""Service Account 认证:SA JSON → JWT → Access Token。"""
|
||||||
from src.core.crypto import crypto_service
|
from src.core.crypto import crypto_service
|
||||||
from src.core.exceptions import InvalidRequestException
|
from src.core.exceptions import InvalidRequestException
|
||||||
@@ -38,11 +38,14 @@ async def _auth_service_account(key: Any) -> ProviderAuthInfo:
|
|||||||
raise InvalidRequestException("Service Account JSON 无效,请重新添加该密钥。")
|
raise InvalidRequestException("Service Account JSON 无效,请重新添加该密钥。")
|
||||||
|
|
||||||
# 获取 Access Token(注入代理配置)
|
# 获取 Access Token(注入代理配置)
|
||||||
|
from src.services.provider.auth import _get_proxy_config
|
||||||
from src.services.proxy_node.resolver import build_proxy_client_kwargs
|
from src.services.proxy_node.resolver import build_proxy_client_kwargs
|
||||||
|
|
||||||
|
effective_proxy = _get_proxy_config(key, endpoint)
|
||||||
|
|
||||||
service = VertexAuthService(sa_json)
|
service = VertexAuthService(sa_json)
|
||||||
access_token = await service.get_access_token(
|
access_token = await service.get_access_token(
|
||||||
httpx_client_kwargs=build_proxy_client_kwargs(timeout=30),
|
httpx_client_kwargs=build_proxy_client_kwargs(effective_proxy, timeout=30),
|
||||||
)
|
)
|
||||||
|
|
||||||
return ProviderAuthInfo(
|
return ProviderAuthInfo(
|
||||||
|
|||||||
@@ -556,7 +556,7 @@ async def get_provider_auth(
|
|||||||
# "vertex_ai" 保留为向后兼容(迁移期间旧数据可能仍使用该值)
|
# "vertex_ai" 保留为向后兼容(迁移期间旧数据可能仍使用该值)
|
||||||
from src.services.provider.adapters.vertex_ai.auth import _auth_service_account
|
from src.services.provider.adapters.vertex_ai.auth import _auth_service_account
|
||||||
|
|
||||||
return await _auth_service_account(key)
|
return await _auth_service_account(key, endpoint)
|
||||||
|
|
||||||
# 标准 API Key:返回 None,由 build_headers 处理
|
# 标准 API Key:返回 None,由 build_headers 处理
|
||||||
return None
|
return None
|
||||||
|
|||||||
126
tests/services/test_provider_auth_vertex_service_account.py
Normal file
126
tests/services/test_provider_auth_vertex_service_account.py
Normal file
@@ -0,0 +1,126 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from src.core.vertex_auth import VertexAuthService
|
||||||
|
from src.services.provider.auth import get_provider_auth
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_provider_auth_vertex_service_account_uses_provider_proxy(
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
sa_json = {
|
||||||
|
"client_email": "svc@example.iam.gserviceaccount.com",
|
||||||
|
"private_key": "-----BEGIN PRIVATE KEY-----\nTEST\n-----END PRIVATE KEY-----\n",
|
||||||
|
"project_id": "demo-project",
|
||||||
|
}
|
||||||
|
provider_proxy = {"node_id": "provider-node", "enabled": True}
|
||||||
|
provider = SimpleNamespace(proxy=provider_proxy)
|
||||||
|
endpoint = SimpleNamespace(provider=provider)
|
||||||
|
key = SimpleNamespace(
|
||||||
|
auth_type="service_account",
|
||||||
|
auth_config="enc_cfg",
|
||||||
|
api_key="enc_key",
|
||||||
|
provider=provider,
|
||||||
|
proxy=None,
|
||||||
|
)
|
||||||
|
|
||||||
|
captured: dict[str, object] = {}
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"src.core.crypto.crypto_service.decrypt",
|
||||||
|
lambda value: json.dumps(sa_json) if value == "enc_cfg" else "",
|
||||||
|
)
|
||||||
|
|
||||||
|
def _fake_build_proxy_client_kwargs(
|
||||||
|
proxy_config: dict[str, object] | None = None,
|
||||||
|
*,
|
||||||
|
timeout: float = 30.0,
|
||||||
|
**_: object,
|
||||||
|
) -> dict[str, object]:
|
||||||
|
captured["proxy_config"] = proxy_config
|
||||||
|
captured["timeout"] = timeout
|
||||||
|
return {"timeout": timeout}
|
||||||
|
|
||||||
|
async def _fake_get_access_token(
|
||||||
|
self: VertexAuthService,
|
||||||
|
*,
|
||||||
|
httpx_client_kwargs: dict[str, object] | None = None,
|
||||||
|
) -> str:
|
||||||
|
captured["httpx_client_kwargs"] = httpx_client_kwargs
|
||||||
|
return "ya29.test-token"
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"src.services.proxy_node.resolver.build_proxy_client_kwargs",
|
||||||
|
_fake_build_proxy_client_kwargs,
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(VertexAuthService, "get_access_token", _fake_get_access_token)
|
||||||
|
|
||||||
|
auth = await get_provider_auth(endpoint, key) # type: ignore[arg-type]
|
||||||
|
|
||||||
|
assert auth is not None
|
||||||
|
assert captured["proxy_config"] == provider_proxy
|
||||||
|
assert captured["timeout"] == 30
|
||||||
|
assert captured["httpx_client_kwargs"] == {"timeout": 30}
|
||||||
|
assert auth.auth_header == "Authorization"
|
||||||
|
assert auth.auth_value == "Bearer ya29.test-token"
|
||||||
|
assert auth.decrypted_auth_config == sa_json
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_provider_auth_vertex_service_account_prefers_key_proxy(
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
sa_json = {
|
||||||
|
"client_email": "svc@example.iam.gserviceaccount.com",
|
||||||
|
"private_key": "-----BEGIN PRIVATE KEY-----\nTEST\n-----END PRIVATE KEY-----\n",
|
||||||
|
"project_id": "demo-project",
|
||||||
|
}
|
||||||
|
provider = SimpleNamespace(proxy={"node_id": "provider-node", "enabled": True})
|
||||||
|
endpoint = SimpleNamespace(provider=provider)
|
||||||
|
key_proxy = {"node_id": "key-node", "enabled": True}
|
||||||
|
key = SimpleNamespace(
|
||||||
|
auth_type="service_account",
|
||||||
|
auth_config="enc_cfg",
|
||||||
|
api_key="enc_key",
|
||||||
|
provider=provider,
|
||||||
|
proxy=key_proxy,
|
||||||
|
)
|
||||||
|
|
||||||
|
captured: dict[str, object] = {}
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"src.core.crypto.crypto_service.decrypt",
|
||||||
|
lambda value: json.dumps(sa_json) if value == "enc_cfg" else "",
|
||||||
|
)
|
||||||
|
|
||||||
|
def _fake_build_proxy_client_kwargs(
|
||||||
|
proxy_config: dict[str, object] | None = None,
|
||||||
|
*,
|
||||||
|
timeout: float = 30.0,
|
||||||
|
**_: object,
|
||||||
|
) -> dict[str, object]:
|
||||||
|
captured["proxy_config"] = proxy_config
|
||||||
|
return {"timeout": timeout}
|
||||||
|
|
||||||
|
async def _fake_get_access_token(
|
||||||
|
self: VertexAuthService,
|
||||||
|
*,
|
||||||
|
httpx_client_kwargs: dict[str, object] | None = None,
|
||||||
|
) -> str:
|
||||||
|
return "ya29.test-token"
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"src.services.proxy_node.resolver.build_proxy_client_kwargs",
|
||||||
|
_fake_build_proxy_client_kwargs,
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(VertexAuthService, "get_access_token", _fake_get_access_token)
|
||||||
|
|
||||||
|
auth = await get_provider_auth(endpoint, key) # type: ignore[arg-type]
|
||||||
|
|
||||||
|
assert auth is not None
|
||||||
|
assert captured["proxy_config"] == key_proxy
|
||||||
43
tests/services/test_vertex_auth.py
Normal file
43
tests/services/test_vertex_auth.py
Normal file
@@ -0,0 +1,43 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from src.core.vertex_auth import VertexAuthError, VertexAuthService
|
||||||
|
|
||||||
|
|
||||||
|
class _TimeoutAsyncClient:
|
||||||
|
def __init__(self, **_: object) -> None:
|
||||||
|
pass
|
||||||
|
|
||||||
|
async def __aenter__(self) -> "_TimeoutAsyncClient":
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> bool:
|
||||||
|
return False
|
||||||
|
|
||||||
|
async def post(self, *args: object, **kwargs: object) -> object:
|
||||||
|
raise httpx.ReadTimeout("")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_vertex_auth_timeout_error_includes_readable_message(
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
service = VertexAuthService(
|
||||||
|
json.dumps(
|
||||||
|
{
|
||||||
|
"client_email": "svc@example.iam.gserviceaccount.com",
|
||||||
|
"private_key": "not-used-in-test",
|
||||||
|
"project_id": "demo-project",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
monkeypatch.setattr(service, "_create_jwt", lambda: "signed-jwt")
|
||||||
|
monkeypatch.setattr("src.core.vertex_auth.httpx.AsyncClient", _TimeoutAsyncClient)
|
||||||
|
|
||||||
|
with pytest.raises(VertexAuthError, match=r"request timed out after 30s"):
|
||||||
|
await service.get_access_token(httpx_client_kwargs={"timeout": 30})
|
||||||
Reference in New Issue
Block a user