tests/test_model_fallback.py

from __future__ import annotations

import pytest

from ergo_agentic.models.routing import (
    ainvoke_structured_with_fallback,
    get_fallback_model_id,
)


class _ServerError(Exception):
    def __init__(self, status_code: int):
        super().__init__(f"status={status_code}")
        self.status_code = status_code


class _StubRunnable:
    def __init__(self, result=None, exc: Exception | None = None):
        self._result = result
        self._exc = exc

    async def ainvoke(self, _messages):
        if self._exc is not None:
            raise self._exc
        return self._result


class _StubModel:
    def __init__(self, runnable: _StubRunnable):
        self._runnable = runnable

    def with_structured_output(self, _schema):
        return self._runnable


@pytest.mark.asyncio
async def test_ainvoke_structured_falls_back_on_transient_error(monkeypatch):
    models = {
        "google_genai:gemini-2.5-pro": _StubModel(_StubRunnable(exc=_ServerError(503))),
        "google_genai:gemini-2.5-flash": _StubModel(_StubRunnable(result={"ok": True})),
    }

    monkeypatch.setattr(
        "ergo_agentic.models.routing.get_chat_model",
        lambda model_id, temperature=0.0: models[model_id],
    )

    result = await ainvoke_structured_with_fallback(
        model_id="google_genai:gemini-2.5-pro",
        schema=dict,
        messages=[],
    )

    assert result == {"ok": True}


@pytest.mark.asyncio
async def test_ainvoke_structured_does_not_fallback_on_non_transient_error(monkeypatch):
    models = {
        "google_genai:gemini-2.5-pro": _StubModel(_StubRunnable(exc=ValueError("bad prompt"))),
    }

    monkeypatch.setattr(
        "ergo_agentic.models.routing.get_chat_model",
        lambda model_id, temperature=0.0: models[model_id],
    )

    with pytest.raises(ValueError, match="bad prompt"):
        await ainvoke_structured_with_fallback(
            model_id="google_genai:gemini-2.5-pro",
            schema=dict,
            messages=[],
        )


def test_get_fallback_model_id_for_gemini_pro():
    assert get_fallback_model_id("google_genai:gemini-2.5-pro") == (
        "google_genai:gemini-2.5-flash"
    )