Aether/tests/services/antigravity/test_client.py

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import AsyncMock, patch

import httpx
import pytest

import src.services.provider.adapters.antigravity.rust_http as antigravity_rust_http_mod
import src.services.request.execution_runtime_client as runtime_client_mod
from src.services.provider.adapters.antigravity.client import (
    fetch_available_models,
    load_code_assist,
    onboard_user,
    parse_retry_delay,
)
from src.services.provider.adapters.antigravity.constants import (
    DAILY_BASE_URL,
    PROD_BASE_URL,
    SANDBOX_BASE_URL,
)
from src.services.request.execution_runtime_client import ExecutionRuntimeSyncResult

# ---------------------------------------------------------------------------
# load_code_assist
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
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=[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=[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"{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)


@pytest.mark.asyncio
async def test_load_code_assist_uses_rust_executor(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(antigravity_rust_http_mod.config, "executor_backend", "rust")
    monkeypatch.setattr(
        runtime_client_mod.ExecutionRuntimeClient,
        "execute_sync_json",
        AsyncMock(
            return_value=ExecutionRuntimeSyncResult(
                status_code=200,
                headers={"content-type": "application/json"},
                response_json={"cloudaicompanionProject": "project-rust"},
            )
        ),
    )
    monkeypatch.setattr(
        "src.clients.http_client.HTTPClientPool.get_proxy_client",
        AsyncMock(side_effect=AssertionError("python fallback should not run")),
    )
    monkeypatch.setattr(
        "src.services.provider.adapters.antigravity.client.url_availability.get_ordered_urls",
        lambda prefer_daily=True: [SANDBOX_BASE_URL],
    )

    data = await load_code_assist("tok", proxy_config=None, timeout_seconds=1.0)

    assert data["cloudaicompanionProject"] == "project-rust"
    plan = runtime_client_mod.ExecutionRuntimeClient.execute_sync_json.await_args.args[0]
    assert plan.url == f"{SANDBOX_BASE_URL}/v1internal:loadCodeAssist"
    assert plan.provider_api_format == "antigravity:load_code_assist"


# ---------------------------------------------------------------------------
# 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)


@pytest.mark.asyncio
async def test_fetch_available_models_uses_rust_executor(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(antigravity_rust_http_mod.config, "executor_backend", "rust")
    monkeypatch.setattr(
        runtime_client_mod.ExecutionRuntimeClient,
        "execute_sync_json",
        AsyncMock(
            return_value=ExecutionRuntimeSyncResult(
                status_code=200,
                headers={"content-type": "application/json"},
                response_json={"models": {"claude-sonnet-4": {"displayName": "Claude Sonnet 4"}}},
            )
        ),
    )
    monkeypatch.setattr(
        "src.clients.http_client.HTTPClientPool.get_proxy_client",
        AsyncMock(side_effect=AssertionError("python fallback should not run")),
    )
    monkeypatch.setattr(
        "src.services.provider.adapters.antigravity.client.url_availability.get_ordered_urls",
        lambda prefer_daily=True: [DAILY_BASE_URL],
    )

    data = await fetch_available_models(
        "tok",
        project_id="project-1",
        proxy_config=None,
        timeout_seconds=1.0,
    )

    assert "models" in data
    plan = runtime_client_mod.ExecutionRuntimeClient.execute_sync_json.await_args.args[0]
    assert plan.url == f"{DAILY_BASE_URL}/v1internal:fetchAvailableModels"
    assert plan.provider_api_format == "antigravity:fetch_available_models"
    assert plan.body.json_body == {"project": "project-1"}


@pytest.mark.asyncio
async def test_onboard_user_uses_rust_executor(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(antigravity_rust_http_mod.config, "executor_backend", "rust")
    monkeypatch.setattr(
        runtime_client_mod.ExecutionRuntimeClient,
        "execute_sync_json",
        AsyncMock(
            return_value=ExecutionRuntimeSyncResult(
                status_code=200,
                headers={"content-type": "application/json"},
                response_json={
                    "done": True,
                    "response": {"cloudaicompanionProject": {"id": "project-onboard"}},
                },
            )
        ),
    )
    monkeypatch.setattr(
        "src.clients.http_client.HTTPClientPool.get_proxy_client",
        AsyncMock(side_effect=AssertionError("python fallback should not run")),
    )
    monkeypatch.setattr(
        "src.services.provider.adapters.antigravity.client.url_availability.get_ordered_urls",
        lambda prefer_daily=True: [PROD_BASE_URL],
    )

    project_id = await onboard_user(
        "tok",
        tier_id="LEGACY",
        proxy_config=None,
        timeout_seconds=1.0,
        max_attempts=1,
    )

    assert project_id == "project-onboard"
    plan = runtime_client_mod.ExecutionRuntimeClient.execute_sync_json.await_args.args[0]
    assert plan.url == f"{PROD_BASE_URL}/v1internal:onboardUser"
    assert plan.provider_api_format == "antigravity:onboard_user"


# ---------------------------------------------------------------------------
# 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