fix(vertex): SA 认证注入代理配置,细化 token 获取异常处理

- _auth_service_account 接收 endpoint 参数,通过 _get_proxy_config 解析代理
- vertex_auth 区分 TimeoutException/RequestError/通用异常,提供可读错误信息
- 新增测试覆盖代理传递和超时场景
This commit is contained in:
fawney19
2026-03-20 00:49:44 +08:00
parent 28fa03451c
commit 772f2ea601
5 changed files with 184 additions and 4 deletions

View 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

View 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})