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