maskanx_cm_backend/tests/test_gemini_vertex_probe.py

# -*- coding: utf-8 -*-
"""The Gemini/Vertex connectivity probe used when a user first adds credentials.

Real incident this guards against: the probe called `generate_content` with
one hardcoded model, gemini-3.6-flash — the newest entry in the registry.
Vertex AI Model Garden availability is scoped per Google Cloud project and
region, so a brand-new model can be enabled in one project's Vertex AI
console before another. Two laptops with byte-identical MaskanX code and
the same Gemini API key produced opposite results: "connected" on one,
"model does not exist or is not available" on the other — because the
*model*, not the key, differed in availability, and the probe reported
the whole connection as broken over exactly one model.
"""
from adclaw.providers.registry import GEMINI_MODELS
from adclaw.providers.store import (
    _GEMINI_PROBE_MODEL_IDS,
    _is_missing_credentials_error,
    _is_model_not_found_error,
    _test_native_gemini_provider_connection,
)

# Verbatim from a real failure: a machine with Vertex enabled but no
# service account JSON and no `gcloud auth application-default login`
# ever run. Raised client-side by google-auth before any network call is
# made — every model would fail identically, since none of them are the
# actual problem.
_MISSING_ADC_ERROR = (
    "Your default credentials were not found. To set up Application "
    "Default Credentials, see "
    "https://cloud.google.com/docs/authentication/external/set-up-adc "
    "for more information."
)


class _FakeModels:
    """Stands in for `client.aio.models`. Queued outcomes, one per call."""

    def __init__(self, outcomes):
        self._outcomes = list(outcomes)
        self.calls = []

    async def generate_content(self, model, contents, config):
        self.calls.append(model)
        outcome = self._outcomes.pop(0)
        if isinstance(outcome, Exception):
            raise outcome
        return outcome


class _FakeAio:
    def __init__(self, models):
        self.models = models


class _FakeClient:
    def __init__(self, outcomes):
        self.aio = _FakeAio(_FakeModels(outcomes))


def _not_found(model_id: str) -> Exception:
    return Exception(
        f"404 NOT_FOUND. {{'error': {{'code': 404, 'message': "
        f"'Publisher Model `{model_id}` not found.'}}}}",
    )


def test_probe_model_ids_exclude_the_customtools_variant():
    """That id is a tool-calling mode of gemini-3.1-pro-preview, not a
    distinct model Vertex serves — probing it separately proves nothing."""
    assert "gemini-3.1-pro-preview-customtools" not in _GEMINI_PROBE_MODEL_IDS
    assert set(_GEMINI_PROBE_MODEL_IDS) == {
        m.id for m in GEMINI_MODELS if "customtools" not in m.id
    }


async def test_a_missing_newest_model_falls_through(monkeypatch):
    """The exact scenario reported: the newest model 404s in this project;
    an older one is fine. The connection must be reported as working."""
    client = _FakeClient([_not_found("gemini-3.6-flash"), "ok"])
    monkeypatch.setattr(
        "adclaw.providers.store._gemini_vertex_enabled", lambda: True,
    )
    monkeypatch.setattr(
        "adclaw.providers.store._create_gemini_vertex_client",
        lambda genai: client,
    )

    result = await _test_native_gemini_provider_connection(
        "Gemini (Google)", "AIzafake",
    )

    assert result["success"] is True
    assert "gemini-3.5-flash" in result["message"]
    # Proves it actually tried the first model and moved on, not that it
    # skipped straight to the second.
    assert client.aio.models.calls[0] == "gemini-3.6-flash"
    assert client.aio.models.calls[1] == "gemini-3.5-flash"


async def test_every_model_missing_is_a_real_failure(monkeypatch):
    """If nothing in the whole registry is available, that is worth
    reporting — the fallback must not paper over a genuinely broken
    project (e.g. Vertex AI API never enabled for it)."""
    outcomes = [_not_found(mid) for mid in _GEMINI_PROBE_MODEL_IDS]
    client = _FakeClient(outcomes)
    monkeypatch.setattr(
        "adclaw.providers.store._gemini_vertex_enabled", lambda: True,
    )
    monkeypatch.setattr(
        "adclaw.providers.store._create_gemini_vertex_client",
        lambda genai: client,
    )

    result = await _test_native_gemini_provider_connection(
        "Gemini (Google)", "AIzafake",
    )

    assert result["success"] is False
    assert "does not exist or is not available" in result["message"]
    assert len(client.aio.models.calls) == len(_GEMINI_PROBE_MODEL_IDS)


async def test_a_real_error_on_the_first_model_stops_immediately(monkeypatch):
    """A bad API key or missing billing applies no matter which model is
    asked for — trying four more models would just be four more of the
    same failure, and delay telling the operator what is actually wrong."""
    client = _FakeClient([Exception("403 PERMISSION_DENIED")])
    monkeypatch.setattr(
        "adclaw.providers.store._gemini_vertex_enabled", lambda: True,
    )
    monkeypatch.setattr(
        "adclaw.providers.store._create_gemini_vertex_client",
        lambda genai: client,
    )

    result = await _test_native_gemini_provider_connection(
        "Gemini (Google)", "AIzafake",
    )

    assert result["success"] is False
    assert "permission" in result["message"].lower()
    assert len(client.aio.models.calls) == 1


async def test_client_construction_failure_is_reported_gracefully(monkeypatch):
    """Regression check: the import + client construction moved outside the
    per-model loop when the fallback was added. If either raises, the
    caller must still get {"success": False, ...}, not an unhandled
    exception — bad service-account JSON must not 500 the endpoint."""
    def _boom(genai):
        raise ValueError("Invalid service account JSON")

    monkeypatch.setattr(
        "adclaw.providers.store._gemini_vertex_enabled", lambda: True,
    )
    monkeypatch.setattr(
        "adclaw.providers.store._create_gemini_vertex_client", _boom,
    )

    result = await _test_native_gemini_provider_connection(
        "Gemini (Google)", "AIzafake",
    )

    assert result["success"] is False
    assert "Invalid service account JSON" in result["message"]


def test_is_model_not_found_error_recognises_graph_style_404s():
    assert _is_model_not_found_error(_not_found("gemini-3.6-flash")) is True
    assert _is_model_not_found_error(Exception("403 PERMISSION_DENIED")) is False


# --- the real incident: missing local Google credentials, not a model problem ---
#
# Reported live: a second laptop, same API key, same real Project ID,
# failed on every one of the five fallback models with the exact same
# "model does not exist" message. The raw error, logged only after this
# fix, was google-auth's own "Your default credentials were not found" —
# a client-side failure before any request reaches Google, misread as a
# per-model 404 because its own text contains "not found".


def test_missing_credentials_error_is_not_mistaken_for_a_missing_model():
    """The bug: this error's own text contains "not found", so the naive
    substring check in _is_model_not_found_error treated "you have no
    Google credentials on this machine" as "this one model is missing" —
    and retried four more times for an identical, unfixable-by-retrying
    failure."""
    exc = Exception(_MISSING_ADC_ERROR)

    assert _is_missing_credentials_error(exc) is True
    assert _is_model_not_found_error(exc) is False


async def test_missing_credentials_fails_fast_without_trying_every_model(
    monkeypatch,
):
    """Retrying five models for a machine-wide auth problem wastes five
    round trips and still arrives at a wrong final message. It must stop
    on the first attempt."""
    client = _FakeClient([Exception(_MISSING_ADC_ERROR)])
    monkeypatch.setattr(
        "adclaw.providers.store._gemini_vertex_enabled", lambda: True,
    )
    monkeypatch.setattr(
        "adclaw.providers.store._create_gemini_vertex_client",
        lambda genai: client,
    )

    result = await _test_native_gemini_provider_connection(
        "Gemini (Google)", "AIzafake",
    )

    assert result["success"] is False
    assert len(client.aio.models.calls) == 1
    assert "credentials" in result["message"].lower()
    # The old, wrong message must not appear — that was the entire bug.
    assert "does not exist or is not available" not in result["message"]


async def test_missing_credentials_message_says_what_to_do_and_that_its_per_machine(
    monkeypatch,
):
    client = _FakeClient([Exception(_MISSING_ADC_ERROR)])
    monkeypatch.setattr(
        "adclaw.providers.store._gemini_vertex_enabled", lambda: True,
    )
    monkeypatch.setattr(
        "adclaw.providers.store._create_gemini_vertex_client",
        lambda genai: client,
    )

    result = await _test_native_gemini_provider_connection(
        "Gemini (Google)", "AIzafake",
    )

    message = result["message"].lower()
    assert "gcloud auth application-default login" in message
    assert "service account" in message
    assert "per-machine" in message or "another computer" in message