/tests/providers/test_vertex.py

"""Tests for Google Vertex AI authentication, discovery, and Chat Completions."""

import asyncio
from unittest.mock import AsyncMock, patch

import httpx
import pytest
from google.auth.credentials import Credentials
from google.auth.exceptions import DefaultCredentialsError, TransportError
from google.auth.transport.requests import Request

from free_claude_code.application.errors import (
    ApplicationUnavailableError,
    InvalidRequestError,
)
from free_claude_code.application.model_metadata import ProviderModelInfo
from free_claude_code.config.provider_catalog import VERTEX_AI_API_ROOT
from free_claude_code.core.failures import ExecutionFailure, FailureKind
from free_claude_code.core.reasoning import ReasoningEffort, ReasoningPolicy
from free_claude_code.providers.base import ProviderConfig
from free_claude_code.providers.model_listing import ModelListResponseError
from free_claude_code.providers.vertex import VertexProvider
from free_claude_code.providers.vertex.auth import GoogleAccessTokenProvider
from free_claude_code.providers.vertex.endpoint import (
    vertex_openai_base_url,
    vertex_publisher_models_url,
    vertex_service_endpoint,
)
from free_claude_code.providers.vertex.models import extract_vertex_model_page
from tests.providers.request_factory import make_messages_request
from tests.providers.support import immediate_admission, reasoning_for

_PROJECT_ID = "my-project"
_GLOBAL_OPENAI_BASE = (
    "https://aiplatform.googleapis.com/v1/projects/my-project/locations/global/"
    "endpoints/openapi"
)
_GLOBAL_MODELS_URL = (
    "https://aiplatform.googleapis.com/v1beta1/publishers/google/models"
)


class FakeCredentials(Credentials):
    """Minimal mutable Google credentials for deterministic refresh tests."""

    def __init__(
        self,
        *,
        token: str | None = None,
        expired: bool = True,
        refresh_error: Exception | None = None,
    ) -> None:
        super().__init__()
        self.token = token
        self._expired = expired
        self._refresh_error = refresh_error
        self.refresh_count = 0
        self.refresh_request: object | None = None

    @property
    def expired(self) -> bool:
        return self._expired

    def refresh(self, request: object) -> None:
        self.refresh_request = request
        self.refresh_count += 1
        if self._refresh_error is not None:
            raise self._refresh_error
        self.token = "refreshed-token"
        self._expired = False


def _token_provider(
    credentials: FakeCredentials | None = None,
) -> GoogleAccessTokenProvider:
    value = credentials or FakeCredentials(token="access-token", expired=False)

    def loader() -> Credentials:
        return value

    return GoogleAccessTokenProvider(loader)


def _provider(
    *,
    location: str = "global",
    token_provider: GoogleAccessTokenProvider | None = None,
) -> VertexProvider:
    return VertexProvider(
        ProviderConfig(api_key="", base_url=VERTEX_AI_API_ROOT),
        project_id=_PROJECT_ID,
        location=location,
        admission=immediate_admission(),
        access_token_provider=token_provider or _token_provider(),
    )


def _simulate_openai_sdk_wire_json(body: dict) -> dict:
    wire = {key: value for key, value in body.items() if key != "extra_body"}
    sdk_extra = body.get("extra_body")
    if isinstance(sdk_extra, dict):
        wire.update(sdk_extra)
    return wire


def _google_thinking_config(wire: dict) -> dict | None:
    literal_extra_body = wire.get("extra_body")
    if not isinstance(literal_extra_body, dict):
        return None
    google = literal_extra_body.get("google")
    if not isinstance(google, dict):
        return None
    thinking_config = google.get("thinking_config")
    return thinking_config if isinstance(thinking_config, dict) else None


@pytest.mark.parametrize(
    ("location", "service_endpoint"),
    [
        ("global", "https://aiplatform.googleapis.com"),
        ("us-central1", "https://us-central1-aiplatform.googleapis.com"),
    ],
)
def test_vertex_endpoints_use_global_or_regional_hosts(
    location: str, service_endpoint: str
) -> None:
    assert vertex_service_endpoint(location) == service_endpoint
    assert vertex_openai_base_url("project/name", location) == (
        f"{service_endpoint}/v1/projects/project%2Fname/locations/{location}/"
        "endpoints/openapi"
    )
    assert vertex_publisher_models_url(location) == (
        f"{service_endpoint}/v1beta1/publishers/google/models"
    )


@pytest.mark.parametrize("project_id", ["", "   "])
def test_vertex_endpoint_requires_project_id(project_id: str) -> None:
    with pytest.raises(ApplicationUnavailableError, match="VERTEX_PROJECT_ID"):
        vertex_openai_base_url(project_id, "global")


@pytest.mark.parametrize("location", ["", "US central1", "us/central1"])
def test_vertex_endpoint_rejects_unsafe_locations(location: str) -> None:
    with pytest.raises(ApplicationUnavailableError, match="VERTEX_LOCATION"):
        vertex_service_endpoint(location)


@pytest.mark.asyncio
async def test_access_token_provider_reuses_valid_token_without_refresh() -> None:
    credentials = FakeCredentials(token="cached-token", expired=False)
    token_provider = _token_provider(credentials)

    assert await token_provider() == "cached-token"
    assert await token_provider() == "cached-token"
    assert credentials.refresh_count == 0


@pytest.mark.asyncio
async def test_access_token_provider_coalesces_concurrent_refreshes() -> None:
    credentials = FakeCredentials()
    token_provider = _token_provider(credentials)

    tokens = await asyncio.gather(*(token_provider() for _ in range(10)))

    assert tokens == ["refreshed-token"] * 10
    assert credentials.refresh_count == 1


@pytest.mark.asyncio
async def test_access_token_refresh_uses_the_vertex_proxy() -> None:
    credentials = FakeCredentials()
    token_provider = GoogleAccessTokenProvider(
        lambda: credentials,
        proxy="socks5://proxy.test:1080",
    )

    assert await token_provider() == "refreshed-token"
    assert isinstance(credentials.refresh_request, Request)
    session = credentials.refresh_request.session
    assert session.proxies == {
        "http": "socks5://proxy.test:1080",
        "https": "socks5://proxy.test:1080",
    }


@pytest.mark.asyncio
async def test_missing_adc_is_non_retryable_authentication_failure() -> None:
    def missing_credentials() -> Credentials:
        raise DefaultCredentialsError("sensitive local path")

    token_provider = GoogleAccessTokenProvider(missing_credentials)

    with pytest.raises(ExecutionFailure) as exc_info:
        await token_provider()

    failure = exc_info.value
    assert failure.kind is FailureKind.AUTHENTICATION
    assert failure.status_code == 401
    assert failure.retryable is False
    assert "gcloud auth application-default login" in failure.message
    assert "sensitive local path" not in failure.message


@pytest.mark.asyncio
async def test_transient_adc_refresh_failure_is_retryable() -> None:
    credentials = FakeCredentials(
        refresh_error=TransportError("temporary auth service failure")
    )
    token_provider = _token_provider(credentials)

    with pytest.raises(ExecutionFailure) as exc_info:
        await token_provider()

    failure = exc_info.value
    assert failure.kind is FailureKind.UNAVAILABLE
    assert failure.status_code == 503
    assert failure.retryable is True
    assert "temporary auth service failure" not in failure.message


def test_vertex_provider_supplies_renewable_token_callback_to_openai() -> None:
    token_provider = _token_provider()
    with (
        patch(
            "free_claude_code.providers.openai_chat.provider.AsyncOpenAI"
        ) as openai_client,
        patch("free_claude_code.providers.vertex.client.httpx.AsyncClient"),
    ):
        provider = _provider(token_provider=token_provider)

    assert provider._provider_name == "VERTEX"
    assert provider._base_url == _GLOBAL_OPENAI_BASE
    assert openai_client.call_args.kwargs["api_key"] is token_provider
    assert openai_client.call_args.kwargs["default_headers"] == {
        "x-goog-user-project": _PROJECT_ID
    }


def test_vertex_request_uses_google_thinking_budget_without_named_effort() -> None:
    provider = _provider()
    request = make_messages_request(
        "google/gemini-3.5-flash",
        thinking={"type": "enabled", "budget_tokens": 2048},
    )

    body = provider._build_request_body(request, reasoning=reasoning_for(request))

    assert body["model"] == "google/gemini-3.5-flash"
    assert "reasoning_effort" not in body
    assert body["extra_body"]["extra_body"]["google"]["thinking_config"] == {
        "include_thoughts": True,
        "thinking_budget": 2048,
    }


def test_vertex_request_maps_reasoning_off_to_zero_budget() -> None:
    provider = _provider()
    request = make_messages_request(
        "google/gemini-3.5-flash",
        thinking={"type": "disabled"},
    )

    body = provider._build_request_body(request, reasoning=reasoning_for(request))

    assert body["extra_body"]["extra_body"]["google"]["thinking_config"] == {
        "thinking_budget": 0,
        "include_thoughts": False,
    }


@pytest.mark.parametrize(
    ("reasoning", "expected_thinking_config"),
    [
        (ReasoningPolicy.provider_default(), None),
        (ReasoningPolicy.off(), {"thinking_budget": 0, "include_thoughts": False}),
        (ReasoningPolicy.on(), {"include_thoughts": True}),
        (
            ReasoningPolicy.on(effort=ReasoningEffort.HIGH),
            {"thinking_budget": 2048, "include_thoughts": True},
        ),
        (
            ReasoningPolicy.on(budget_tokens=777),
            {"thinking_budget": 777, "include_thoughts": True},
        ),
        (
            ReasoningPolicy.on(
                effort=ReasoningEffort.HIGH,
                budget_tokens=777,
            ),
            {"thinking_budget": 777, "include_thoughts": True},
        ),
    ],
)
def test_vertex_reasoning_has_one_google_wire_owner(
    reasoning: ReasoningPolicy,
    expected_thinking_config: dict | None,
) -> None:
    provider = _provider()

    body = provider._build_request_body(
        make_messages_request("google/gemini", thinking=None),
        reasoning=reasoning,
    )
    wire = _simulate_openai_sdk_wire_json(body)

    assert "reasoning_effort" not in wire
    assert _google_thinking_config(wire) == expected_thinking_config


def test_vertex_preserves_caller_thinking_config_only_for_provider_default() -> None:
    provider = _provider()
    request = make_messages_request(
        "google/gemini",
        thinking=None,
        extra_body={
            "extra_body": {
                "google": {
                    "thinking_config": {
                        "thinking_level": "low",
                        "include_thoughts": False,
                    }
                }
            }
        },
    )

    body = provider._build_request_body(
        request,
        reasoning=ReasoningPolicy.provider_default(),
    )
    wire = _simulate_openai_sdk_wire_json(body)

    assert _google_thinking_config(wire) == {
        "thinking_level": "low",
        "include_thoughts": False,
    }


def test_vertex_rejects_caller_thinking_config_with_fcc_reasoning_control() -> None:
    provider = _provider()
    request = make_messages_request(
        "google/gemini",
        thinking=None,
        extra_body={
            "extra_body": {"google": {"thinking_config": {"thinking_level": "low"}}}
        },
    )

    with pytest.raises(InvalidRequestError, match="thinking_config"):
        provider._build_request_body(
            request,
            reasoning=ReasoningPolicy.on(effort=ReasoningEffort.HIGH),
        )


def test_vertex_model_page_translates_google_resource_names_generically() -> None:
    model_ids, page_token = extract_vertex_model_page(
        {
            "publisherModels": [
                {"name": "publishers/google/models/gemini-3.5-flash"},
                {"name": "publishers/acme/models/custom-chat"},
            ],
            "nextPageToken": "next-page",
        }
    )

    assert model_ids == frozenset({"google/gemini-3.5-flash", "acme/custom-chat"})
    assert page_token == "next-page"


@pytest.mark.parametrize(
    "payload",
    [
        None,
        {},
        {"publisherModels": "not-a-list"},
        {"publisherModels": [{}]},
        {"publisherModels": [{"name": "models/missing-publisher"}]},
        {"publisherModels": [{"name": "publishers/google/models/   "}]},
        {"publisherModels": [], "nextPageToken": 123},
    ],
)
def test_vertex_model_page_rejects_malformed_responses(payload: object) -> None:
    with pytest.raises(ModelListResponseError, match="VERTEX model-list response"):
        extract_vertex_model_page(payload)


@pytest.mark.asyncio
async def test_vertex_model_discovery_follows_native_pagination() -> None:
    provider = _provider()
    responses = [
        httpx.Response(
            200,
            json={
                "publisherModels": [
                    {"name": "publishers/google/models/gemini-3.5-flash"}
                ],
                "nextPageToken": "page-2",
            },
            request=httpx.Request("GET", _GLOBAL_MODELS_URL),
        ),
        httpx.Response(
            200,
            json={
                "publisherModels": [{"name": "publishers/google/models/gemini-3.1-pro"}]
            },
            request=httpx.Request("GET", _GLOBAL_MODELS_URL),
        ),
    ]
    with patch.object(
        provider._model_list_client,
        "get",
        new_callable=AsyncMock,
        side_effect=responses,
    ) as get:
        model_infos = await provider.list_model_infos()

    assert model_infos == frozenset(
        {
            ProviderModelInfo("google/gemini-3.5-flash"),
            ProviderModelInfo("google/gemini-3.1-pro"),
        }
    )
    assert get.await_args_list[0].kwargs == {
        "params": None,
        "headers": {
            "Authorization": "Bearer access-token",
            "x-goog-user-project": _PROJECT_ID,
        },
    }
    assert get.await_args_list[1].kwargs == {
        "params": {"pageToken": "page-2"},
        "headers": {
            "Authorization": "Bearer access-token",
            "x-goog-user-project": _PROJECT_ID,
        },
    }
    assert all(response.is_closed for response in responses)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "response",
    [
        httpx.Response(
            200,
            json={"publisherModels": []},
            request=httpx.Request("GET", _GLOBAL_MODELS_URL),
        ),
        httpx.Response(
            200,
            content=b"not-json",
            request=httpx.Request("GET", _GLOBAL_MODELS_URL),
        ),
    ],
)
async def test_vertex_model_discovery_rejects_unusable_success_response(
    response: httpx.Response,
) -> None:
    provider = _provider()
    with (
        patch.object(
            provider._model_list_client,
            "get",
            new_callable=AsyncMock,
            return_value=response,
        ),
        pytest.raises(ModelListResponseError, match="VERTEX model-list response"),
    ):
        await provider.list_model_infos()

    assert response.is_closed


@pytest.mark.asyncio
async def test_vertex_model_discovery_rejects_repeated_page_token() -> None:
    provider = _provider()
    responses = [
        httpx.Response(
            200,
            json={"publisherModels": [], "nextPageToken": "same"},
            request=httpx.Request("GET", _GLOBAL_MODELS_URL),
        ),
        httpx.Response(
            200,
            json={"publisherModels": [], "nextPageToken": "same"},
            request=httpx.Request("GET", _GLOBAL_MODELS_URL),
        ),
    ]
    with (
        patch.object(
            provider._model_list_client,
            "get",
            new_callable=AsyncMock,
            side_effect=responses,
        ),
        pytest.raises(ModelListResponseError, match="repeated nextPageToken"),
    ):
        await provider.list_model_infos()