awoooi/apps/api/tests/test_alert_rule_ai_provider_route.py

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock

import pytest

from src.services import ai_control as ai_control_module
from src.services import ai_rate_limiter as limiter_module
from src.services import alert_rule_engine as rule_module


@pytest.fixture(autouse=True)
def _enable_explicit_gemini_env(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("ENABLE_GEMINI", "true")


class _Response:
    def raise_for_status(self) -> None:
        return None

    def json(self) -> dict:
        return {"response": "generated-rule"}


class _OllamaClient:
    def __init__(self) -> None:
        self.urls: list[str] = []

    async def __aenter__(self) -> _OllamaClient:
        return self

    async def __aexit__(self, *_args: object) -> None:
        return None

    async def post(self, url: str, **_kwargs: object) -> _Response:
        self.urls.append(url)
        if len(self.urls) < 3:
            raise RuntimeError("isolated Ollama failure")
        return _Response()


@pytest.mark.asyncio
async def test_background_rule_generation_ignores_custom_route_insertion(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    endpoints = [
        SimpleNamespace(url="http://gcp-a", provider_name="ollama_gcp_a"),
        SimpleNamespace(url="http://gcp-b", provider_name="ollama_gcp_b"),
        SimpleNamespace(url="http://host111", provider_name="ollama_local"),
    ]
    client = _OllamaClient()
    monkeypatch.setattr(rule_module, "resolve_ollama_order", lambda _task: endpoints)
    monkeypatch.setattr(rule_module.httpx, "AsyncClient", lambda **_kwargs: client)

    result = await rule_module._call_ollama(
        "generate rule",
        "http://untrusted-custom-route",
        "test-model",
    )

    assert result == "generated-rule"
    assert client.urls == [
        "http://gcp-a/api/generate",
        "http://gcp-b/api/generate",
        "http://host111/api/generate",
    ]


@pytest.mark.asyncio
async def test_background_rule_generation_never_readds_disabled_gemini(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    limiter = MagicMock()
    limiter.reserve_generation = AsyncMock()
    monkeypatch.setattr(
        ai_control_module,
        "is_provider_disabled",
        AsyncMock(return_value=True),
    )
    monkeypatch.setattr(limiter_module, "get_ai_rate_limiter", lambda: limiter)

    result = await rule_module._call_gemini(
        "generate rule",
        "configured-test-value",
        context={"trace_id": "trace-1"},
    )

    assert result is None
    limiter.reserve_generation.assert_not_awaited()


@pytest.mark.asyncio
async def test_background_rule_generation_stops_before_http_when_cost_guard_blocks(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    limiter = MagicMock()
    limiter.reserve_generation = AsyncMock(
        return_value=SimpleNamespace(
            allowed=False,
            receipt_id="receipt-blocked",
            reason="total_cost_limit",
        )
    )
    monkeypatch.setattr(
        ai_control_module,
        "is_provider_disabled",
        AsyncMock(return_value=False),
    )
    monkeypatch.setattr(limiter_module, "get_ai_rate_limiter", lambda: limiter)

    result = await rule_module._call_gemini(
        "generate rule",
        "configured-test-value",
        context={
            "trace_id": "trace-1",
            "run_id": "run-1",
            "work_item_id": "work-1",
        },
    )

    assert result is None
    limiter.reserve_generation.assert_awaited_once()


@pytest.mark.asyncio
async def test_background_rule_generation_bills_thought_tokens_as_output(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    response = MagicMock()
    response.raise_for_status.return_value = None
    response.json.return_value = {
        "candidates": [{"content": {"parts": [{"text": "generated"}]}}],
        "usageMetadata": {
            "promptTokenCount": 3,
            "candidatesTokenCount": 5,
            "thoughtsTokenCount": 7,
        },
    }
    client = MagicMock()
    client.__aenter__ = AsyncMock(return_value=client)
    client.__aexit__ = AsyncMock(return_value=None)
    client.post = AsyncMock(return_value=response)
    reservation = SimpleNamespace(
        allowed=True,
        receipt_id="receipt-thoughts",
        reason="allowed",
    )
    limiter = MagicMock()
    limiter.reserve_generation = AsyncMock(return_value=reservation)
    limiter.finalize_generation = AsyncMock(return_value=True)
    monkeypatch.setattr(
        ai_control_module,
        "is_provider_disabled",
        AsyncMock(return_value=False),
    )
    monkeypatch.setattr(limiter_module, "get_ai_rate_limiter", lambda: limiter)
    monkeypatch.setattr(rule_module.httpx, "AsyncClient", lambda **_kwargs: client)

    result = await rule_module._call_gemini(
        "generate rule",
        "configured-test-value",
        context={
            "trace_id": "trace-thoughts",
            "run_id": "run-thoughts",
            "work_item_id": "work-thoughts",
        },
    )

    assert result == "generated"
    limiter.finalize_generation.assert_awaited_once_with(
        reservation,
        success=True,
        prompt_tokens=3,
        completion_tokens=12,
    )