947 lines
36 KiB
Python
947 lines
36 KiB
Python
"""agy provider tests — OAuth paste-URL login, translation, pool rotation.
|
|
|
|
Everything external is faked at the ``httpx.AsyncClient`` seam, so the suite
|
|
runs with no Google, no network and no accounts. What is *not* faked is the
|
|
protocol shape: the fake answers are the shapes Google answers with (token
|
|
JSON, userinfo JSON, ``v1internal:generateContent`` responses, 429s), so a
|
|
regression in our translation fails here rather than in production.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import time
|
|
from urllib.parse import parse_qs, urlparse
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
|
|
import app.infrastructure.adapters.agy as agy
|
|
from app.application.ports import ChatMessage
|
|
from app.core.errors import ProviderUnavailableError
|
|
from app.domain.ai_control.models import AiModel, ModelCredential
|
|
from tests.conftest import auth_header, signup
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fake httpx seam
|
|
# ---------------------------------------------------------------------------
|
|
|
|
GOOD_GENERATE: dict = {
|
|
"candidates": [{"content": {"parts": [{"text": "سلام"}]}, "finishReason": "STOP"}],
|
|
"usageMetadata": {
|
|
"promptTokenCount": 7,
|
|
"candidatesTokenCount": 3,
|
|
"thoughtsTokenCount": 0,
|
|
"totalTokenCount": 10,
|
|
},
|
|
"model": "gemini-3-flash",
|
|
}
|
|
|
|
|
|
class _FakeResponse:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
status_code: int = 200,
|
|
json_data: dict | None = None,
|
|
text: str = "",
|
|
headers: dict | None = None,
|
|
) -> None:
|
|
self.status_code = status_code
|
|
self._json = json_data if json_data is not None else {}
|
|
self._text = text
|
|
self.headers = headers or {}
|
|
|
|
@property
|
|
def is_success(self) -> bool:
|
|
return self.status_code < 400
|
|
|
|
def json(self) -> dict:
|
|
return self._json
|
|
|
|
@property
|
|
def text(self) -> str:
|
|
return self._text
|
|
|
|
|
|
def _token_response(refresh: str = "refresh-1", access: str = "access-1") -> _FakeResponse:
|
|
return _FakeResponse(
|
|
json_data={"access_token": access, "refresh_token": refresh, "expires_in": 3600}
|
|
)
|
|
|
|
|
|
def patch_google(monkeypatch, handler) -> list[dict]:
|
|
"""Point ``httpx.AsyncClient`` at a scripted handler.
|
|
|
|
The adapter creates its own ``AsyncClient`` internally, so the seam is the
|
|
class itself. ``handler(method, url, json_body, headers, data)`` returns a
|
|
``_FakeResponse`` (sync or async); every call is recorded and returned for
|
|
assertions.
|
|
"""
|
|
import inspect
|
|
|
|
calls: list[dict] = []
|
|
|
|
async def _dispatch(method, url, json_body, headers, data):
|
|
result = handler(method, url, json_body, headers or {}, data or {})
|
|
if inspect.isawaitable(result):
|
|
result = await result
|
|
return result
|
|
|
|
async def _post(self, url, *, json=None, headers=None, data=None, **_):
|
|
calls.append({"method": "POST", "url": url, "json": json, "headers": headers, "data": data})
|
|
return await _dispatch("POST", url, json, headers, data)
|
|
|
|
async def _get(self, url, *, headers=None, **_):
|
|
calls.append({"method": "GET", "url": url, "json": None, "headers": headers, "data": None})
|
|
return await _dispatch("GET", url, None, headers, None)
|
|
|
|
class _Client:
|
|
def __init__(self, *_, **__):
|
|
pass
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *_) -> bool:
|
|
return False
|
|
|
|
post = _post
|
|
get = _get
|
|
|
|
monkeypatch.setattr(agy.httpx, "AsyncClient", _Client)
|
|
return calls
|
|
|
|
|
|
def default_handler(method, url, json_body, headers, data):
|
|
"""The Google every happy-path test talks to."""
|
|
if url.endswith("/token"):
|
|
return _token_response()
|
|
if "userinfo" in url:
|
|
return _FakeResponse(json_data={"email": "person@gmail.com"})
|
|
if url.endswith("/v1internal:loadCodeAssist"):
|
|
return _FakeResponse(json_data={"cloudaicompanionProject": "agy-proj-123"})
|
|
if url.endswith("/v1internal:fetchAvailableModels"):
|
|
return _FakeResponse(json_data={"models": {"gemini-3-flash": {}, "claude-sonnet-4-5": {}}})
|
|
if url.endswith("/v1internal:generateContent"):
|
|
return _FakeResponse(json_data=GOOD_GENERATE)
|
|
raise AssertionError(f"unexpected {method} {url}")
|
|
|
|
|
|
@pytest.fixture
|
|
def google(monkeypatch):
|
|
"""Async handler matching ``default_handler``; returns the call log."""
|
|
return patch_google(monkeypatch, lambda *a: default_handler(*a))
|
|
|
|
|
|
@pytest.fixture
|
|
def decrypted(monkeypatch):
|
|
"""Make ``decrypt`` a passthrough so tests can build fake ciphertext."""
|
|
monkeypatch.setattr(agy, "decrypt", lambda blob: blob.split("PLAINTEXT::", 1)[1])
|
|
|
|
|
|
def _agy_credential(
|
|
*, email: str, refresh_token: str = "refresh-1", project_id: str = "agy-proj-123"
|
|
) -> ModelCredential:
|
|
"""A credential as the adapter sees it: ciphertext already holds the blob."""
|
|
envelope = json.dumps(
|
|
{"refresh_token": refresh_token, "project_id": project_id, "email": email}
|
|
)
|
|
return ModelCredential(
|
|
provider_id="prov-1",
|
|
name=email,
|
|
encrypted_api_key=f"PLAINTEXT::{envelope}",
|
|
email=email,
|
|
project_id=project_id,
|
|
)
|
|
|
|
|
|
_MODEL = AiModel(provider_id="p", internal_name="gemini-3-flash", display_name_for_admin="f")
|
|
_MESSAGES = [ChatMessage(role="user", content="سلام")]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# HTTP endpoints (platform owner only)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_agy_endpoints_require_platform_owner(client, google):
|
|
body = await signup(client, "customer@example.com")
|
|
response = await client.post(
|
|
"/api/admin/agy/accounts/login/start", headers=auth_header(body["access_token"])
|
|
)
|
|
assert response.status_code == 403
|
|
|
|
|
|
async def test_agy_login_flow_end_to_end(client, google, platform_owner_token):
|
|
# 1. start the headless login.
|
|
started = await client.post(
|
|
"/api/admin/agy/accounts/login/start",
|
|
headers=auth_header(platform_owner_token),
|
|
)
|
|
assert started.status_code == 201, started.text
|
|
login_id = started.json()["login_id"]
|
|
auth_url = started.json()["auth_url"]
|
|
state = parse_qs(urlparse(auth_url).query)["state"][0]
|
|
|
|
# 2. the operator pastes the browser's callback URL back.
|
|
completed = await client.post(
|
|
"/api/admin/agy/accounts/login/complete",
|
|
headers=auth_header(platform_owner_token),
|
|
json={
|
|
"login_id": login_id,
|
|
"callback_url": (f"http://localhost:51121/oauth-callback?code=auth-code&state={state}"),
|
|
},
|
|
)
|
|
assert completed.status_code == 201, completed.text
|
|
account = completed.json()
|
|
assert account["email"] == "person@gmail.com"
|
|
assert account["status"] == "active"
|
|
assert account["is_active"] is True
|
|
|
|
# No credential material ever crosses the API surface.
|
|
assert "refresh_token" not in completed.text
|
|
assert "encrypted" not in completed.text
|
|
assert "agy-proj-123" in completed.text # project id is admin metadata
|
|
|
|
# 3. the account is listed.
|
|
listed = await client.get("/api/admin/agy/accounts", headers=auth_header(platform_owner_token))
|
|
assert listed.status_code == 200
|
|
emails = [a["email"] for a in listed.json()]
|
|
assert emails == ["person@gmail.com"]
|
|
|
|
|
|
async def test_agy_complete_login_rejects_unrelated_url(client, google, platform_owner_token):
|
|
started = await client.post(
|
|
"/api/admin/agy/accounts/login/start",
|
|
headers=auth_header(platform_owner_token),
|
|
)
|
|
completed = await client.post(
|
|
"/api/admin/agy/accounts/login/complete",
|
|
headers=auth_header(platform_owner_token),
|
|
json={
|
|
"login_id": started.json()["login_id"],
|
|
"callback_url": "http://localhost:51121/oauth-callback?code=x&state=elsewhere",
|
|
},
|
|
)
|
|
assert completed.status_code == 422
|
|
# The failed attempt must not have created an account.
|
|
listed = await client.get("/api/admin/agy/accounts", headers=auth_header(platform_owner_token))
|
|
assert listed.json() == []
|
|
|
|
|
|
async def test_agy_account_test_and_delete(client, google, platform_owner_token):
|
|
started = await client.post(
|
|
"/api/admin/agy/accounts/login/start",
|
|
headers=auth_header(platform_owner_token),
|
|
)
|
|
state = parse_qs(urlparse(started.json()["auth_url"]).query)["state"][0]
|
|
created = await client.post(
|
|
"/api/admin/agy/accounts/login/complete",
|
|
headers=auth_header(platform_owner_token),
|
|
json={
|
|
"login_id": started.json()["login_id"],
|
|
"callback_url": (f"http://localhost:51121/oauth-callback?code=c&state={state}"),
|
|
},
|
|
)
|
|
account_id = created.json()["id"]
|
|
|
|
# A test call runs a real (faked-Google) completion through the account.
|
|
tested = await client.post(
|
|
f"/api/admin/agy/accounts/{account_id}/test",
|
|
headers=auth_header(platform_owner_token),
|
|
)
|
|
assert tested.status_code == 200, tested.text
|
|
assert tested.json()["ok"] is True
|
|
assert tested.json()["detail"] == "سلام"
|
|
|
|
deleted = await client.delete(
|
|
f"/api/admin/agy/accounts/{account_id}",
|
|
headers=auth_header(platform_owner_token),
|
|
)
|
|
assert deleted.status_code == 204
|
|
listed = await client.get("/api/admin/agy/accounts", headers=auth_header(platform_owner_token))
|
|
assert listed.json() == []
|
|
|
|
|
|
async def test_agy_same_account_twice_replaces(client, google, platform_owner_token):
|
|
"""Re-login of a known Google account rotates the key, never duplicates."""
|
|
for _ in range(2):
|
|
started = await client.post(
|
|
"/api/admin/agy/accounts/login/start",
|
|
headers=auth_header(platform_owner_token),
|
|
)
|
|
state = parse_qs(urlparse(started.json()["auth_url"]).query)["state"][0]
|
|
done = await client.post(
|
|
"/api/admin/agy/accounts/login/complete",
|
|
headers=auth_header(platform_owner_token),
|
|
json={
|
|
"login_id": started.json()["login_id"],
|
|
"callback_url": (f"http://localhost:51121/oauth-callback?code=c&state={state}"),
|
|
},
|
|
)
|
|
assert done.status_code == 201
|
|
listed = await client.get("/api/admin/agy/accounts", headers=auth_header(platform_owner_token))
|
|
assert len(listed.json()) == 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Seed: the agy catalog
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def seeded_install(monkeypatch, engine, session_factory):
|
|
"""Run the real ``app.seed`` against the in-memory test database."""
|
|
import app.seed as seed_module
|
|
|
|
monkeypatch.setattr(seed_module, "engine", engine)
|
|
monkeypatch.setattr(seed_module, "SessionFactory", session_factory)
|
|
await seed_module.seed()
|
|
return seed_module
|
|
|
|
|
|
async def _agy_models(session_factory) -> list[tuple[str, bool]]:
|
|
from sqlalchemy import select
|
|
|
|
from app.infrastructure.db.models import AiModelRow, AiProviderRow
|
|
|
|
async with session_factory() as session:
|
|
rows = (
|
|
await session.execute(
|
|
select(AiModelRow, AiProviderRow.name)
|
|
.join(AiProviderRow, AiModelRow.provider_id == AiProviderRow.id)
|
|
.where(AiProviderRow.name == "agy")
|
|
)
|
|
).all()
|
|
return [(row.internal_name, row.is_active) for row, _name in rows]
|
|
|
|
|
|
async def test_seed_creates_agy_catalog(client, seeded_install, session_factory):
|
|
models = await _agy_models(session_factory)
|
|
names = {name for name, _ in models}
|
|
assert names == {
|
|
"gemini-3-flash",
|
|
"gemini-3-pro-high",
|
|
"claude-sonnet-4-5",
|
|
"claude-sonnet-4-5-thinking",
|
|
"claude-opus-4-5-thinking",
|
|
"gpt-oss-120b-medium",
|
|
}
|
|
# agy is the only provider: the catalog is routable from the first run.
|
|
assert all(active is True for _, active in models)
|
|
|
|
|
|
async def test_seed_round_trips_efforts(client, seeded_install, session_factory):
|
|
"""The seed's effort column is what the adapter later sends to Google."""
|
|
from sqlalchemy import select
|
|
|
|
from app.infrastructure.db.models import AiModelRow, AiProviderRow
|
|
|
|
async with session_factory() as session:
|
|
rows = (
|
|
await session.execute(
|
|
select(AiModelRow)
|
|
.join(AiProviderRow, AiModelRow.provider_id == AiProviderRow.id)
|
|
.where(AiProviderRow.name == "agy")
|
|
)
|
|
).scalars()
|
|
efforts = {row.internal_name: row.reasoning_effort for row in rows}
|
|
assert efforts["gemini-3-pro-high"] == "high"
|
|
assert efforts["claude-sonnet-4-5-thinking"] == "high"
|
|
assert efforts["gemini-3-flash"] == ""
|
|
|
|
|
|
async def test_seed_leaves_no_legacy_providers(client, seeded_install, session_factory):
|
|
"""agy is the only provider row a seeded install may carry."""
|
|
from sqlalchemy import select
|
|
|
|
from app.infrastructure.db.models import AiProviderRow
|
|
|
|
async with session_factory() as session:
|
|
rows = (await session.execute(select(AiProviderRow))).scalars()
|
|
types = {row.type for row in rows}
|
|
names = {row.name for row in rows}
|
|
assert types == {"agy"}
|
|
assert "mock" not in names
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# OAuth login flow
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_auth_url_shape():
|
|
challenge = agy.AgyAuthClient().begin()
|
|
query = parse_qs(urlparse(challenge.auth_url).query)
|
|
assert challenge.auth_url.startswith("https://accounts.google.com/o/oauth2/v2/auth")
|
|
assert query["client_id"][0] == (
|
|
"1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
|
|
)
|
|
assert query["redirect_uri"][0] == "http://localhost:51121/oauth-callback"
|
|
assert query["code_challenge_method"][0] == "S256"
|
|
# access_type=offline + prompt=consent is what makes Google hand back a
|
|
# refresh token; without it the login completes but nothing can be stored.
|
|
assert query["access_type"][0] == "offline"
|
|
assert query["prompt"][0] == "consent"
|
|
assert "cloud-platform" in query["scope"][0]
|
|
assert challenge.verifier and challenge.state
|
|
|
|
|
|
async def test_paste_flow_exchanges_and_resolves_account(google):
|
|
"""start → paste callback URL → account material, verifier stays server-side."""
|
|
login_id, _auth_url = await agy.begin_login()
|
|
challenge = agy._CHALLENGES[login_id]
|
|
callback = f"http://localhost:51121/oauth-callback?code=auth-code&state={challenge.state}"
|
|
result = await agy.redeem_login(login_id, callback)
|
|
assert result.email == "person@gmail.com"
|
|
assert result.project_id == "agy-proj-123"
|
|
assert result.refresh_token == "refresh-1"
|
|
# The challenge is single-use.
|
|
with pytest.raises(agy.AgyOAuthError):
|
|
await agy.redeem_login(login_id, callback)
|
|
|
|
|
|
async def test_foreign_state_is_rejected(google):
|
|
login_id, _ = await agy.begin_login()
|
|
with pytest.raises(agy.AgyOAuthError):
|
|
await agy.redeem_login(
|
|
login_id, "http://localhost:51121/oauth-callback?code=x&state=not-ours"
|
|
)
|
|
|
|
|
|
async def test_missing_code_is_rejected(google):
|
|
login_id, _ = await agy.begin_login()
|
|
with pytest.raises(agy.AgyOAuthError):
|
|
await agy.redeem_login(login_id, "http://localhost:51121/oauth-callback?state=x")
|
|
|
|
|
|
async def test_google_error_page_is_rejected(google):
|
|
login_id, _ = await agy.begin_login()
|
|
with pytest.raises(agy.AgyOAuthError):
|
|
await agy.redeem_login(
|
|
login_id, "http://localhost:51121/oauth-callback?error=access_denied"
|
|
)
|
|
|
|
|
|
async def test_refresh_token_missing_is_actionable(monkeypatch):
|
|
"""A token response without ``refresh_token`` fails loudly and tells why."""
|
|
|
|
async def handler(method, url, json_body, headers, data):
|
|
if url.endswith("/token"):
|
|
return _FakeResponse(json_data={"access_token": "a", "expires_in": 3600})
|
|
return await default_handler(method, url, json_body, headers, data)
|
|
|
|
patch_google(monkeypatch, handler)
|
|
login_id, _ = await agy.begin_login()
|
|
challenge = agy._CHALLENGES[login_id]
|
|
with pytest.raises(agy.AgyOAuthError) as excinfo:
|
|
await agy.redeem_login(
|
|
login_id,
|
|
f"http://localhost:51121/oauth-callback?code=c&state={challenge.state}",
|
|
)
|
|
assert "refresh_token" in str(excinfo.value)
|
|
|
|
|
|
async def test_rejected_refresh_is_an_account_error_not_an_outage(monkeypatch):
|
|
async def handler(method, url, json_body, headers, data):
|
|
if url.endswith("/token"):
|
|
return _FakeResponse(
|
|
status_code=400,
|
|
text=json.dumps({"error": "invalid_grant"}),
|
|
)
|
|
return await default_handler(method, url, json_body, headers, data)
|
|
|
|
patch_google(monkeypatch, handler)
|
|
with pytest.raises(Exception) as excinfo:
|
|
await agy.AgyAuthClient().refresh("dead-token")
|
|
# Raised as the internal taxonomy so the adapter benches the account.
|
|
assert type(excinfo.value).__name__ == "_AgyAuthError"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Request translation (pure functions)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_payload_translation(decrypted):
|
|
account = agy.AgyAccount(
|
|
credential_id="c1", email="a@b.c", project_id="proj-1", refresh_token="refresh-1"
|
|
)
|
|
payload = agy._build_generate_payload(
|
|
account,
|
|
"gemini-3-pro-high",
|
|
[
|
|
ChatMessage(role="system", content="تو مدیر هستی."),
|
|
ChatMessage(role="user", content="سلام"),
|
|
ChatMessage(role="assistant", content="سلام! چیزی کمکی میخواهید؟"),
|
|
ChatMessage(role="user", content="یک گزارش بساز"),
|
|
],
|
|
system="پرامپت سیستم",
|
|
temperature=0.4,
|
|
max_tokens=512,
|
|
)
|
|
assert payload["model"] == "gemini-3-pro-high"
|
|
assert payload["project"] == "proj-1"
|
|
assert payload["requestId"].startswith("agent-")
|
|
assert payload["request"]["sessionId"].startswith("-")
|
|
# System content from both the `system` argument and system messages lands
|
|
# in systemInstruction, not as a user turn.
|
|
instruction = payload["request"]["systemInstruction"]["parts"][0]["text"]
|
|
assert "پرامپت سیستم" in instruction
|
|
assert "تو مدیر هستی." in instruction
|
|
roles = [c["role"] for c in payload["request"]["contents"]]
|
|
assert roles == ["user", "model", "user"]
|
|
assert payload["request"]["generationConfig"]["temperature"] == 0.4
|
|
assert payload["request"]["generationConfig"]["maxOutputTokens"] == 512
|
|
# Function calling must be VALIDATED or tool schemas get rejected.
|
|
assert payload["request"]["toolConfig"]["functionCallingConfig"]["mode"] == "VALIDATED"
|
|
assert all(s["threshold"] == "OFF" for s in payload["request"]["safetySettings"])
|
|
|
|
|
|
def test_effort_becomes_a_thinking_config_on_a_thinking_model():
|
|
"""Effort is translated, not forwarded: high → a 16384-token budget."""
|
|
account = agy.AgyAccount(
|
|
credential_id="c1", email="a@b.c", project_id="proj-1", refresh_token="refresh-1"
|
|
)
|
|
payload = agy._build_generate_payload(
|
|
account,
|
|
"claude-sonnet-4-5-thinking",
|
|
[ChatMessage(role="user", content="سلام")],
|
|
system="",
|
|
temperature=0.4,
|
|
max_tokens=512,
|
|
reasoning_effort="high",
|
|
)
|
|
generation = payload["request"]["generationConfig"]
|
|
assert generation["thinkingConfig"] == {
|
|
"thinkingBudget": 16384,
|
|
"includeThoughts": True,
|
|
}
|
|
# The budget must fit inside the answer window.
|
|
assert generation["maxOutputTokens"] == 16384 + 8000
|
|
|
|
|
|
def test_effort_is_not_sent_to_a_non_thinking_model():
|
|
"""``thinkingConfig`` on a model that cannot think is a hard 400."""
|
|
account = agy.AgyAccount(
|
|
credential_id="c1", email="a@b.c", project_id="proj-1", refresh_token="refresh-1"
|
|
)
|
|
payload = agy._build_generate_payload(
|
|
account,
|
|
"claude-sonnet-4-5",
|
|
[ChatMessage(role="user", content="سلام")],
|
|
system="",
|
|
temperature=0.4,
|
|
max_tokens=512,
|
|
reasoning_effort="high",
|
|
)
|
|
assert "thinkingConfig" not in payload["request"]["generationConfig"]
|
|
assert payload["request"]["generationConfig"]["maxOutputTokens"] == 512
|
|
|
|
|
|
def test_unknown_effort_falls_back_to_medium():
|
|
account = agy.AgyAccount(
|
|
credential_id="c1", email="a@b.c", project_id="proj-1", refresh_token="refresh-1"
|
|
)
|
|
payload = agy._build_generate_payload(
|
|
account,
|
|
"claude-opus-4-5-thinking",
|
|
[ChatMessage(role="user", content="سلام")],
|
|
system="",
|
|
temperature=0.4,
|
|
max_tokens=512,
|
|
reasoning_effort="turbo",
|
|
)
|
|
budget = payload["request"]["generationConfig"]["thinkingConfig"]["thinkingBudget"]
|
|
assert budget == 8192
|
|
|
|
|
|
def test_response_parsing_skips_thoughts():
|
|
model = AiModel(provider_id="p", internal_name="m", display_name_for_admin="m")
|
|
response = agy._parse_generate_response(
|
|
{
|
|
"candidates": [
|
|
{
|
|
"content": {
|
|
"parts": [
|
|
{"text": "دارم فکر میکنم", "thought": True},
|
|
{"text": "پاسخ نهایی"},
|
|
]
|
|
},
|
|
"finishReason": "STOP",
|
|
}
|
|
],
|
|
"usageMetadata": {
|
|
"promptTokenCount": 11,
|
|
"candidatesTokenCount": 5,
|
|
"thoughtsTokenCount": 4,
|
|
},
|
|
},
|
|
model=model,
|
|
started=0.0,
|
|
)
|
|
assert response.content == "پاسخ نهایی"
|
|
assert response.input_tokens == 11
|
|
assert response.output_tokens == 5
|
|
assert response.raw["thought_tokens"] == 4
|
|
assert response.raw["finish_reason"] == "STOP"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# The adapter: completion + account pool rotation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_complete_happy_path(google, decrypted):
|
|
adapter = agy.AgyModelAdapter()
|
|
response = await adapter.complete(
|
|
model=_MODEL,
|
|
credentials=[_agy_credential(email="one@gmail.com")],
|
|
messages=_MESSAGES,
|
|
)
|
|
assert response.content == "سلام"
|
|
assert response.input_tokens == 7
|
|
assert response.output_tokens == 3
|
|
generate = [c for c in google if c["url"].endswith("generateContent")]
|
|
assert len(generate) == 1
|
|
# The account's project id is what bills the call.
|
|
assert generate[0]["json"]["project"] == "agy-proj-123"
|
|
|
|
|
|
async def test_rate_limited_account_is_rotated_and_benched(monkeypatch, decrypted):
|
|
"""A 429 on one account must not fail the request: the pool moves on."""
|
|
|
|
async def handler(method, url, json_body, headers, data):
|
|
if url.endswith("/token"):
|
|
# Each account's refresh token maps to its own access token.
|
|
access = {"refresh-1": "access-one", "refresh-2": "access-two"}[
|
|
data.get("refresh_token", "")
|
|
]
|
|
return _token_response(refresh=data.get("refresh_token", ""), access=access)
|
|
if url.endswith("/v1internal:generateContent"):
|
|
auth = headers.get("Authorization", "")
|
|
if auth.endswith("access-one"):
|
|
# Both base URLs hit the same quota.
|
|
return _FakeResponse(status_code=429, text="quota", headers={"Retry-After": "120"})
|
|
return _FakeResponse(json_data=GOOD_GENERATE)
|
|
return await default_handler(method, url, json_body, headers, data)
|
|
|
|
patch_google(monkeypatch, handler)
|
|
reported: list[tuple] = []
|
|
|
|
async def report(credential_id, *, status=None, cooldown_until=None, error=None, used=False):
|
|
reported.append((credential_id, status, cooldown_until, error, used))
|
|
|
|
adapter = agy.AgyModelAdapter(report=report)
|
|
# Start the rotation cursor so `one` is genuinely attempted first.
|
|
adapter._state.cursor = -1
|
|
one = _agy_credential(email="one@gmail.com", refresh_token="refresh-1")
|
|
two = _agy_credential(email="two@gmail.com", refresh_token="refresh-2")
|
|
response = await adapter.complete(model=_MODEL, credentials=[one, two], messages=_MESSAGES)
|
|
assert response.content == "سلام"
|
|
|
|
cooldowns = [r for r in reported if r[1] == "cooldown"]
|
|
assert cooldowns, "the rate-limited account must be reported"
|
|
assert cooldowns[0][0] == one.id
|
|
# Google's Retry-After is honoured: the bench is ~120s out, not the
|
|
# configured fallback.
|
|
remaining = cooldowns[0][2].timestamp() - time.time()
|
|
assert 90 <= remaining <= 200
|
|
|
|
|
|
async def test_every_account_rate_limited_raises(monkeypatch, decrypted):
|
|
async def handler(method, url, json_body, headers, data):
|
|
if url.endswith("/token"):
|
|
return _token_response()
|
|
if url.endswith("/v1internal:generateContent"):
|
|
return _FakeResponse(status_code=429, text="quota")
|
|
return await default_handler(method, url, json_body, headers, data)
|
|
|
|
patch_google(monkeypatch, handler)
|
|
adapter = agy.AgyModelAdapter()
|
|
with pytest.raises(ProviderUnavailableError):
|
|
await adapter.complete(
|
|
model=_MODEL,
|
|
credentials=[
|
|
_agy_credential(email="one@gmail.com", refresh_token="r1"),
|
|
_agy_credential(email="two@gmail.com", refresh_token="r2"),
|
|
],
|
|
messages=_MESSAGES,
|
|
)
|
|
|
|
|
|
async def test_rejected_grant_marks_account_error(monkeypatch, decrypted):
|
|
"""401 means the OAuth grant died — bench as error, move to next account."""
|
|
|
|
async def handler(method, url, json_body, headers, data):
|
|
if url.endswith("/token"):
|
|
access = {"refresh-1": "access-one", "refresh-2": "access-two"}[
|
|
data.get("refresh_token", "")
|
|
]
|
|
return _token_response(refresh=data.get("refresh_token", ""), access=access)
|
|
if url.endswith("/v1internal:generateContent"):
|
|
auth = headers.get("Authorization", "")
|
|
if auth.endswith("access-one"):
|
|
return _FakeResponse(status_code=401, text="unauthorized")
|
|
return _FakeResponse(json_data=GOOD_GENERATE)
|
|
return await default_handler(method, url, json_body, headers, data)
|
|
|
|
patch_google(monkeypatch, handler)
|
|
reported: list[tuple] = []
|
|
|
|
async def report(credential_id, *, status=None, cooldown_until=None, error=None, used=False):
|
|
reported.append((credential_id, status))
|
|
|
|
adapter = agy.AgyModelAdapter(report=report)
|
|
adapter._state.cursor = -1
|
|
dead = _agy_credential(email="dead@gmail.com", refresh_token="refresh-1")
|
|
alive = _agy_credential(email="alive@gmail.com", refresh_token="refresh-2")
|
|
response = await adapter.complete(model=_MODEL, credentials=[dead, alive], messages=_MESSAGES)
|
|
assert response.content == "سلام"
|
|
errors = [r for r in reported if r[1] == "error"]
|
|
assert errors and errors[0][0] == dead.id
|
|
# The healthy account is cleared from any bench and marked used.
|
|
assert any(r[0] == alive.id and r[1] == "active" for r in reported)
|
|
|
|
|
|
async def test_server_error_falls_through_to_next_backend(monkeypatch, decrypted):
|
|
"""5xx on one backend URL is not the account's fault — try the other."""
|
|
|
|
async def handler(method, url, json_body, headers, data):
|
|
if url.endswith("/token"):
|
|
return _token_response()
|
|
if url.endswith("/v1internal:generateContent"):
|
|
if "daily-cloudcode-pa" in url:
|
|
return _FakeResponse(status_code=503, text="backend down")
|
|
return _FakeResponse(json_data=GOOD_GENERATE)
|
|
return await default_handler(method, url, json_body, headers, data)
|
|
|
|
patch_google(monkeypatch, handler)
|
|
adapter = agy.AgyModelAdapter()
|
|
response = await adapter.complete(
|
|
model=_MODEL,
|
|
credentials=[_agy_credential(email="one@gmail.com")],
|
|
messages=_MESSAGES,
|
|
)
|
|
assert response.content == "سلام"
|
|
|
|
|
|
async def test_no_accounts_raises_actionable_error(google, decrypted):
|
|
adapter = agy.AgyModelAdapter()
|
|
with pytest.raises(ProviderUnavailableError) as excinfo:
|
|
await adapter.complete(model=_MODEL, credentials=[], messages=_MESSAGES)
|
|
assert "اکانت agy" in str(excinfo.value)
|
|
|
|
|
|
async def test_health_reports_account_count(google, decrypted):
|
|
async def loader():
|
|
return [_agy_credential(email="one@gmail.com")]
|
|
|
|
adapter = agy.AgyModelAdapter(loader=loader)
|
|
health = await adapter.health()
|
|
assert health.ok and "1" in health.detail
|
|
|
|
|
|
async def test_available_models(google, decrypted):
|
|
async def loader():
|
|
return [_agy_credential(email="one@gmail.com")]
|
|
|
|
adapter = agy.AgyModelAdapter(loader=loader)
|
|
models = await adapter.available_models()
|
|
assert models == ["claude-sonnet-4-5", "gemini-3-flash"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Live per-account model catalog (the console's model dropdown source)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_available_models_for_one_account(google, decrypted):
|
|
"""The console asks for a specific account, so the fetch is per-account."""
|
|
adapter = agy.AgyModelAdapter()
|
|
models = await adapter.available_models_for(_agy_credential(email="one@gmail.com"))
|
|
assert [m.id for m in models] == ["claude-sonnet-4-5", "gemini-3-flash"]
|
|
assert models[1].supported_efforts == ["low", "medium", "high"]
|
|
# The list is fetched from Google every call — no cached catalog.
|
|
fetches = [c for c in google if c["url"].endswith("fetchAvailableModels")]
|
|
assert len(fetches) == 1
|
|
|
|
|
|
async def test_available_models_for_account_without_material():
|
|
"""A credential with no refresh token can never list models."""
|
|
adapter = agy.AgyModelAdapter()
|
|
credential = ModelCredential(
|
|
provider_id="p",
|
|
name="empty@gmail.com",
|
|
encrypted_api_key="PLAINTEXT::{}",
|
|
email="empty@gmail.com",
|
|
)
|
|
with pytest.raises(ProviderUnavailableError):
|
|
await adapter.available_models_for(credential)
|
|
|
|
|
|
async def test_account_models_endpoint(client, google, platform_owner_token):
|
|
"""GET /accounts/{id}/models returns that account's live catalog."""
|
|
started = await client.post(
|
|
"/api/admin/agy/accounts/login/start",
|
|
headers=auth_header(platform_owner_token),
|
|
)
|
|
state = parse_qs(urlparse(started.json()["auth_url"]).query)["state"][0]
|
|
account_id = (
|
|
await client.post(
|
|
"/api/admin/agy/accounts/login/complete",
|
|
headers=auth_header(platform_owner_token),
|
|
json={
|
|
"login_id": started.json()["login_id"],
|
|
"callback_url": (f"http://localhost:51121/oauth-callback?code=c&state={state}"),
|
|
},
|
|
)
|
|
).json()["id"]
|
|
|
|
response = await client.get(
|
|
f"/api/admin/agy/accounts/{account_id}/models",
|
|
headers=auth_header(platform_owner_token),
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
data = response.json()
|
|
assert [m["id"] for m in data] == ["claude-sonnet-4-5", "gemini-3-flash"]
|
|
assert data[1]["supported_efforts"] == ["low", "medium", "high"]
|
|
|
|
missing = await client.get(
|
|
"/api/admin/agy/accounts/does-not-exist/models",
|
|
headers=auth_header(platform_owner_token),
|
|
)
|
|
assert missing.status_code == 404
|
|
|
|
|
|
async def test_completion_carries_the_model_effort(monkeypatch, decrypted):
|
|
"""The end-to-end path: effort on the model row reaches the payload."""
|
|
|
|
async def handler(method, url, json_body, headers, data):
|
|
if url.endswith("/token"):
|
|
return _token_response()
|
|
if url.endswith("/v1internal:generateContent"):
|
|
return _FakeResponse(json_data=GOOD_GENERATE)
|
|
return await default_handler(method, url, json_body, headers, data)
|
|
|
|
calls = patch_google(monkeypatch, handler)
|
|
model = AiModel(
|
|
provider_id="p",
|
|
internal_name="claude-sonnet-4-5-thinking",
|
|
display_name_for_admin="f",
|
|
reasoning_effort="low",
|
|
)
|
|
adapter = agy.AgyModelAdapter()
|
|
await adapter.complete(
|
|
model=model,
|
|
credentials=[_agy_credential(email="one@gmail.com")],
|
|
messages=_MESSAGES,
|
|
)
|
|
generate = [c for c in calls if c["url"].endswith("generateContent")]
|
|
assert generate[0]["json"]["request"]["generationConfig"]["thinkingConfig"] == {
|
|
"thinkingBudget": 4096,
|
|
"includeThoughts": True,
|
|
}
|
|
|
|
|
|
def test_parse_cli_models_output():
|
|
"""Verify parsing raw output from agy models CLI."""
|
|
sample_output = (
|
|
"Fetching available models...\n"
|
|
"gemini-3.8-flash-high\tGemini 3.8 Flash (High)\n"
|
|
"gemini-3.8-flash-medium\tGemini 3.8 Flash (Medium)\n"
|
|
"gemini-3.8-flash-low\tGemini 3.8 Flash (Low)\n"
|
|
"claude-sonnet-4-6\tClaude Sonnet 4.6 (Thinking)\n"
|
|
"claude-opus-4-6-thinking\tClaude Opus 4.6 (Thinking)\n"
|
|
"gpt-oss-120b-medium\tGPT-OSS 120B (Medium)\n"
|
|
)
|
|
models = agy.AgyModelAdapter._parse_cli_models_output(sample_output)
|
|
assert len(models) == 6
|
|
|
|
m_high = next(m for m in models if m.id == "gemini-3.8-flash-high")
|
|
assert m_high.display_name == "Gemini 3.8 Flash (High)"
|
|
assert m_high.supports_thinking is True
|
|
assert m_high.supported_efforts == ["high"]
|
|
assert m_high.default_effort == "high"
|
|
|
|
m_med = next(m for m in models if m.id == "gemini-3.8-flash-medium")
|
|
assert m_med.supported_efforts == ["medium"]
|
|
assert m_med.default_effort == "medium"
|
|
|
|
m_low = next(m for m in models if m.id == "gemini-3.8-flash-low")
|
|
assert m_low.supported_efforts == ["low"]
|
|
assert m_low.default_effort == "low"
|
|
|
|
m_claude = next(m for m in models if m.id == "claude-sonnet-4-6")
|
|
assert m_claude.supports_thinking is True
|
|
|
|
m_gpt = next(m for m in models if m.id == "gpt-oss-120b-medium")
|
|
assert m_gpt.supported_efforts == ["medium"]
|
|
|
|
|
|
def test_build_generate_payload_tiered_mapping():
|
|
"""Verify virtual CLI models map to underlying tiered model with appropriate budget."""
|
|
account = agy.AgyAccount(
|
|
credential_id="cred-1",
|
|
refresh_token="rt-1",
|
|
email="test@example.com",
|
|
project_id="proj-123",
|
|
)
|
|
messages = [ChatMessage(role="user", content="Hello")]
|
|
|
|
payload_low = agy._build_generate_payload(
|
|
account,
|
|
"gemini-3.8-flash-low",
|
|
messages,
|
|
system="",
|
|
temperature=0.7,
|
|
max_tokens=4096,
|
|
)
|
|
assert payload_low["model"] == "gemini-3.8-flash-tiered"
|
|
assert payload_low["request"]["generationConfig"]["thinkingConfig"] == {
|
|
"thinkingBudget": 4096,
|
|
"includeThoughts": True,
|
|
}
|
|
|
|
payload_high = agy._build_generate_payload(
|
|
account,
|
|
"gemini-3.8-flash-high",
|
|
messages,
|
|
system="",
|
|
temperature=0.7,
|
|
max_tokens=4096,
|
|
)
|
|
assert payload_high["model"] == "gemini-3.8-flash-tiered"
|
|
assert payload_high["request"]["generationConfig"]["thinkingConfig"] == {
|
|
"thinkingBudget": 16384,
|
|
"includeThoughts": True,
|
|
}
|
|
|
|
|
|
async def test_fetch_models_via_cli_mocked(monkeypatch):
|
|
"""Verify _fetch_models_via_cli invokes the CLI binary and handles output."""
|
|
account = agy.AgyAccount(
|
|
credential_id="cred-1",
|
|
refresh_token="rt-test-123",
|
|
email="test@example.com",
|
|
project_id="proj-123",
|
|
)
|
|
|
|
class FakeProc:
|
|
returncode = 0
|
|
|
|
async def communicate(self):
|
|
return b"gemini-3.8-flash-high\tGemini 3.8 Flash (High)\n", b""
|
|
|
|
async def fake_create_subprocess_exec(*args, **kwargs):
|
|
return FakeProc()
|
|
|
|
monkeypatch.setattr(agy.asyncio, "create_subprocess_exec", fake_create_subprocess_exec)
|
|
monkeypatch.setattr(agy.shutil, "which", lambda cmd: "/usr/bin/agy")
|
|
|
|
models = await agy.AgyModelAdapter._fetch_models_via_cli(account)
|
|
assert len(models) == 1
|
|
assert models[0].id == "gemini-3.8-flash-high"
|
|
assert models[0].display_name == "Gemini 3.8 Flash (High)"
|