leocrm/tests/test_cost_protection.py

"""Tests for LLM cost overrun protection (B.17).

Covers:
- Cost cap blocks LLM calls when budget exceeded
- Cost tracking in Redis is correct
- Alerts at 50%/80%/100% of budget
- Hard stop blocks LLM calls
- Reset at month boundary (key includes month)
"""

from __future__ import annotations

import uuid
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from app.ai.llm_client import (
    _get_cost_key,
    get_tenant_monthly_cost,
    _check_tenant_budget,
    _track_tenant_cost,
    _check_cost_alerts,
)


# ──────────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────────

TENANT_ID = uuid.uuid4()


def _mock_redis_with_cost(initial_cost: float = 0.0):
    """Create a mock Redis client with cost tracking support."""
    r = MagicMock()
    r.get = AsyncMock(return_value=str(initial_cost) if initial_cost > 0 else None)
    r.incrbyfloat = AsyncMock(return_value=initial_cost)
    r.expire = AsyncMock()
    r.set = AsyncMock(return_value=True)  # nx=True returns True if key was set
    return r


def _mock_redis_with_alerted(alerted_thresholds: set[float] | None = None):
    """Create a mock Redis where some alert thresholds are already set."""
    r = MagicMock()
    r.get = AsyncMock(return_value=None)
    r.incrbyfloat = AsyncMock(return_value=0.0)
    r.expire = AsyncMock()
    alerted = alerted_thresholds or set()

    async def _set_nx(key, value, nx=True, ex=None):
        # Extract threshold from key: cost:alerted:{tenant_id}:{threshold}:{month}
        parts = key.split(":")
        if len(parts) >= 4:
            try:
                threshold = float(parts[3])
                if threshold in alerted:
                    return None  # Key already exists
            except ValueError:
                pass
        return True  # Key was set (new alert)

    r.set = AsyncMock(side_effect=_set_nx)
    return r


# ──────────────────────────────────────────────────────────────────────────
# Tests
# ──────────────────────────────────────────────────────────────────────────


class TestCostKey:
    """Tests for Redis cost key generation."""

    def test_cost_key_format(self):
        key = _get_cost_key(TENANT_ID)
        now = datetime.now(timezone.utc)
        month_str = now.strftime("%Y-%m")
        assert key == f"cost:tenant:{TENANT_ID}:month:{month_str}"

    def test_cost_key_includes_month(self):
        key = _get_cost_key(TENANT_ID)
        now = datetime.now(timezone.utc)
        assert now.strftime("%Y-%m") in key

    def test_cost_key_with_string_tenant(self):
        key = _get_cost_key("some-tenant-id")
        assert "some-tenant-id" in key


class TestGetTenantMonthlyCost:
    """Tests for get_tenant_monthly_cost."""

    @pytest.mark.asyncio
    async def test_get_cost_returns_zero_when_no_data(self):
        with patch("app.core.auth.get_redis") as mock_get_redis:
            r = MagicMock()
            r.get = AsyncMock(return_value=None)
            mock_get_redis.return_value = r
            cost = await get_tenant_monthly_cost(TENANT_ID)
            assert cost == 0.0

    @pytest.mark.asyncio
    async def test_get_cost_returns_stored_value(self):
        with patch("app.core.auth.get_redis") as mock_get_redis:
            r = MagicMock()
            r.get = AsyncMock(return_value="42.50")
            mock_get_redis.return_value = r
            cost = await get_tenant_monthly_cost(TENANT_ID)
            assert cost == 42.50

    @pytest.mark.asyncio
    async def test_get_cost_handles_redis_error(self):
        with patch("app.core.auth.get_redis", side_effect=Exception("Redis down")):
            cost = await get_tenant_monthly_cost(TENANT_ID)
            assert cost == 0.0


class TestCheckTenantBudget:
    """Tests for budget enforcement."""

    @pytest.mark.asyncio
    async def test_budget_check_passes_when_under_budget(self):
        with patch("app.core.auth.get_redis") as mock_get_redis:
            r = MagicMock()
            r.get = AsyncMock(return_value="10.0")
            mock_get_redis.return_value = r
            with patch("app.config.get_settings") as mock_settings:
                mock_settings.return_value.llm_monthly_budget_usd = 100.0
                mock_settings.return_value.llm_hard_cutoff = True
                await _check_tenant_budget(TENANT_ID)

    @pytest.mark.asyncio
    async def test_budget_check_raises_when_exceeded_hard_cutoff(self):
        with patch("app.core.auth.get_redis") as mock_get_redis:
            r = MagicMock()
            r.get = AsyncMock(return_value="100.0")
            mock_get_redis.return_value = r
            with patch("app.config.get_settings") as mock_settings:
                mock_settings.return_value.llm_monthly_budget_usd = 100.0
                mock_settings.return_value.llm_hard_cutoff = True
                with pytest.raises(ValueError, match="Tenant LLM budget exceeded"):
                    await _check_tenant_budget(TENANT_ID)

    @pytest.mark.asyncio
    async def test_budget_check_no_raise_when_soft_cutoff(self):
        with patch("app.core.auth.get_redis") as mock_get_redis:
            r = MagicMock()
            r.get = AsyncMock(return_value="150.0")
            mock_get_redis.return_value = r
            with patch("app.config.get_settings") as mock_settings:
                mock_settings.return_value.llm_monthly_budget_usd = 100.0
                mock_settings.return_value.llm_hard_cutoff = False
                await _check_tenant_budget(TENANT_ID)

    @pytest.mark.asyncio
    async def test_budget_check_skips_when_no_tenant(self):
        await _check_tenant_budget(None)

    @pytest.mark.asyncio
    async def test_budget_check_skips_when_budget_zero(self):
        with patch("app.config.get_settings") as mock_settings:
            mock_settings.return_value.llm_monthly_budget_usd = 0.0
            mock_settings.return_value.llm_hard_cutoff = True
            await _check_tenant_budget(TENANT_ID)


class TestTrackTenantCost:
    """Tests for cost tracking in Redis."""

    @pytest.mark.asyncio
    async def test_track_cost_increments_redis(self):
        with patch("app.core.auth.get_redis") as mock_get_redis:
            r = MagicMock()
            r.incrbyfloat = AsyncMock(return_value=5.50)
            r.expire = AsyncMock()
            mock_get_redis.return_value = r

            with patch("app.ai.llm_client._check_cost_alerts", new_callable=AsyncMock):
                total = await _track_tenant_cost(TENANT_ID, 5.50)
                assert total == 5.50
                r.incrbyfloat.assert_called_once()
                r.expire.assert_called_once()

    @pytest.mark.asyncio
    async def test_track_cost_skips_zero(self):
        total = await _track_tenant_cost(TENANT_ID, 0.0)
        assert total == 0.0

    @pytest.mark.asyncio
    async def test_track_cost_skips_none_tenant(self):
        total = await _track_tenant_cost(None, 5.0)
        assert total == 0.0

    @pytest.mark.asyncio
    async def test_track_cost_handles_redis_error(self):
        with patch("app.core.auth.get_redis", side_effect=Exception("Redis down")):
            total = await _track_tenant_cost(TENANT_ID, 5.0)
            assert total == 0.0


class TestCostAlerts:
    """Tests for cost alert thresholds."""

    @pytest.mark.asyncio
    async def test_alert_at_50_percent(self):
        with patch("app.core.auth.get_redis") as mock_get_redis:
            r = _mock_redis_with_alerted()
            mock_get_redis.return_value = r
            with patch("app.config.get_settings") as mock_settings:
                mock_settings.return_value.llm_monthly_budget_usd = 100.0
                await _check_cost_alerts(TENANT_ID, 50.0, db=None)
                assert r.set.called

    @pytest.mark.asyncio
    async def test_alert_at_80_percent(self):
        with patch("app.core.auth.get_redis") as mock_get_redis:
            r = _mock_redis_with_alerted()
            mock_get_redis.return_value = r
            with patch("app.config.get_settings") as mock_settings:
                mock_settings.return_value.llm_monthly_budget_usd = 100.0
                await _check_cost_alerts(TENANT_ID, 80.0, db=None)
                assert r.set.called

    @pytest.mark.asyncio
    async def test_alert_at_100_percent(self):
        with patch("app.core.auth.get_redis") as mock_get_redis:
            r = _mock_redis_with_alerted()
            mock_get_redis.return_value = r
            with patch("app.config.get_settings") as mock_settings:
                mock_settings.return_value.llm_monthly_budget_usd = 100.0
                await _check_cost_alerts(TENANT_ID, 100.0, db=None)
                assert r.set.called

    @pytest.mark.asyncio
    async def test_alert_not_fired_below_threshold(self):
        with patch("app.core.auth.get_redis") as mock_get_redis:
            r = _mock_redis_with_alerted()
            mock_get_redis.return_value = r
            with patch("app.config.get_settings") as mock_settings:
                mock_settings.return_value.llm_monthly_budget_usd = 100.0
                await _check_cost_alerts(TENANT_ID, 30.0, db=None)
                r.set.assert_not_called()

    @pytest.mark.asyncio
    async def test_alert_handles_redis_error(self):
        with patch("app.core.auth.get_redis", side_effect=Exception("Redis down")):
            with patch("app.config.get_settings") as mock_settings:
                mock_settings.return_value.llm_monthly_budget_usd = 100.0
                await _check_cost_alerts(TENANT_ID, 50.0, db=None)


class TestMonthReset:
    """Tests for monthly cost reset via key expiration."""

    def test_cost_key_changes_per_month(self):
        key_jan = f"cost:tenant:{TENANT_ID}:month:2026-01"
        key_feb = f"cost:tenant:{TENANT_ID}:month:2026-02"
        assert key_jan != key_feb

    @pytest.mark.asyncio
    async def test_track_cost_sets_ttl(self):
        with patch("app.core.auth.get_redis") as mock_get_redis:
            r = MagicMock()
            r.incrbyfloat = AsyncMock(return_value=10.0)
            r.expire = AsyncMock()
            mock_get_redis.return_value = r
            with patch("app.ai.llm_client._check_cost_alerts", new_callable=AsyncMock):
                await _track_tenant_cost(TENANT_ID, 10.0)
                r.expire.assert_called_once()
                call_args = r.expire.call_args
                ttl = call_args[0][1] if len(call_args[0]) > 1 else call_args[1].get("time")
                assert ttl == 35 * 24 * 3600