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_model_not_found_error,
    _test_native_gemini_provider_connection,
)


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