awoooi/apps/api/tests/test_ollama_provider_endpoints.py

from __future__ import annotations

from typing import Any

import pytest

from src.services.ai_providers import ollama as ollama_module
from src.services.ai_providers.ollama import OllamaGcpBProvider, OllamaProvider


class _FakeRegistry:
    def get_model(self, provider: str, use_case: str) -> str:
        return "qwen3:14b"

    def get_provider_options(self, provider: str) -> dict[str, Any]:
        return {"num_predict": 32, "temperature": 0.1, "top_p": 0.9}


class _FakeResponse:
    status_code = 200

    def raise_for_status(self) -> None:
        return None

    def json(self) -> dict[str, Any]:
        return {
            "response": '{"summary":"ok"}',
            "eval_count": 4,
            "prompt_eval_count": 3,
        }


class _FakeClient:
    def __init__(self) -> None:
        self.posted_urls: list[str] = []
        self.posted_payloads: list[dict[str, Any]] = []
        self.checked_urls: list[str] = []

    async def post(self, url: str, **kwargs: Any) -> _FakeResponse:
        self.posted_urls.append(url)
        self.posted_payloads.append(kwargs.get("json", {}))
        return _FakeResponse()

    async def get(self, url: str, **kwargs: Any) -> _FakeResponse:
        self.checked_urls.append(url)
        return _FakeResponse()


@pytest.mark.asyncio
async def test_ollama_gcp_b_analyze_uses_secondary_url(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(ollama_module, "get_model_registry", lambda: _FakeRegistry())
    monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
    monkeypatch.setattr(
        ollama_module.settings,
        "OLLAMA_SECONDARY_URL",
        "http://secondary:11436",
    )
    monkeypatch.setattr(ollama_module.settings, "ALERT_OLLAMA_MODEL", "qwen3:14b")

    client = _FakeClient()
    provider = OllamaGcpBProvider()

    async def _get_client() -> _FakeClient:
        return client

    monkeypatch.setattr(provider, "_get_client", _get_client)

    result = await provider.analyze("diagnose", context={"task_type": "diagnose"})

    assert result.success is True
    assert result.provider == "ollama_gcp_b"
    assert client.posted_urls == ["http://secondary:11436/api/generate"]
    assert client.posted_payloads[0]["model"] == "qwen3:14b"


@pytest.mark.asyncio
async def test_ollama_gcp_a_allows_heavy_diagnose_model(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(ollama_module, "get_model_registry", lambda: _FakeRegistry())
    monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
    monkeypatch.setattr(ollama_module.settings, "OLLAMA_SECONDARY_URL", "http://secondary:11436")
    monkeypatch.setattr(ollama_module.settings, "ALERT_OLLAMA_MODEL", "qwen3:14b")

    client = _FakeClient()
    provider = OllamaProvider()

    async def _get_client() -> _FakeClient:
        return client

    monkeypatch.setattr(provider, "_get_client", _get_client)

    result = await provider.analyze("diagnose", context={"task_type": "diagnose"})

    assert result.success is True
    assert client.posted_urls == ["http://primary:11435/api/generate"]
    assert client.posted_payloads[0]["model"] == "qwen3:14b"


@pytest.mark.asyncio
async def test_ollama_gcp_a_coerces_non_diagnosis_heavy_model_to_health_model(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(ollama_module, "get_model_registry", lambda: _FakeRegistry())
    monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
    monkeypatch.setattr(ollama_module.settings, "OLLAMA_SECONDARY_URL", "http://secondary:11436")
    monkeypatch.setattr(ollama_module.settings, "OLLAMA_HEALTH_CHECK_MODEL", "gemma3:4b")

    client = _FakeClient()
    provider = OllamaProvider()

    async def _get_client() -> _FakeClient:
        return client

    monkeypatch.setattr(provider, "_get_client", _get_client)

    result = await provider.analyze("background summary", context={"task_type": "summary"})

    assert result.success is True
    assert client.posted_urls == ["http://primary:11435/api/generate"]
    assert client.posted_payloads[0]["model"] == "gemma3:4b"


@pytest.mark.asyncio
async def test_ollama_gcp_a_can_explicitly_allow_heavy_model(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(ollama_module, "get_model_registry", lambda: _FakeRegistry())
    monkeypatch.setattr(ollama_module.settings, "OLLAMA_URL", "http://primary:11435")
    monkeypatch.setattr(ollama_module.settings, "OLLAMA_SECONDARY_URL", "http://secondary:11436")
    monkeypatch.setattr(ollama_module.settings, "ALERT_OLLAMA_MODEL", "qwen3:14b")

    client = _FakeClient()
    provider = OllamaProvider()

    async def _get_client() -> _FakeClient:
        return client

    monkeypatch.setattr(provider, "_get_client", _get_client)

    result = await provider.analyze(
        "deep diagnose",
        context={"task_type": "diagnose", "allow_gcp_heavy_model": True},
    )

    assert result.success is True
    assert client.posted_payloads[0]["model"] == "qwen3:14b"


@pytest.mark.asyncio
async def test_ollama_base_provider_health_uses_endpoint_hook(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    client = _FakeClient()
    provider = OllamaProvider()
    monkeypatch.setattr(provider, "_endpoint_url", lambda: "http://primary:11435")

    async def _get_client() -> _FakeClient:
        return client

    monkeypatch.setattr(provider, "_get_client", _get_client)

    assert await provider.health_check() is True
    assert client.checked_urls == ["http://primary:11435/api/tags"]