litellm/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py

"""
Tests for the /guardrails/usage/* endpoints backing the dashboard Guardrail Monitor.

Regression (LIT-2529): guardrails defined in config.yaml live only in
IN_MEMORY_GUARDRAIL_HANDLER, so the monitor's overview/detail/logs endpoints —
which read the litellm_guardrailstable Prisma table — could not see them:
detail 404'd, overview omitted them (or rendered them as Custom/Guardrail
orphans), and logs missed their logical-name alias.
"""

from datetime import datetime
from typing import Any, Optional
from unittest.mock import AsyncMock, MagicMock, patch

import pytest


from fastapi import HTTPException
from prisma.errors import TableNotFoundError

from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler
from litellm.proxy.guardrails.usage_endpoints import (
    guardrails_usage_detail,
    guardrails_usage_logs,
    guardrails_usage_overview,
    policies_usage_overview,
)
from litellm.types.guardrails import Guardrail, LitellmParams

ADMIN = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
# Query() defaults don't resolve to None when the handler is called directly.
START, END = "2026-04-20", "2026-04-27"


def _config_handler(*guardrails: Guardrail) -> InMemoryGuardrailHandler:
    """A real handler seeded with config-sourced YAML guardrails (no callbacks)."""
    handler = InMemoryGuardrailHandler()
    for g in guardrails:
        gid = g["guardrail_id"]
        handler.IN_MEMORY_GUARDRAILS[gid] = g
        handler._sources[gid] = "config"
    return handler


def _yaml_guardrail(
    guardrail_id: str = "yaml-1",
    name: str = "yaml-pii",
    provider: str = "presidio",
    info: Optional[dict] = None,
) -> Guardrail:
    return Guardrail(
        guardrail_id=guardrail_id,
        guardrail_name=name,
        litellm_params=LitellmParams(guardrail=provider, mode="pre_call"),
        guardrail_info=info if info is not None else {"type": "PII", "description": "yaml-defined"},
    )


def _db_row(guardrail_id: str = "db-1", name: str = "db-guard", provider: str = "aim") -> Any:
    """A Prisma-style row: attribute access, litellm_params/guardrail_info as plain dicts."""
    row = MagicMock(spec=["guardrail_id", "guardrail_name", "litellm_params", "guardrail_info"])
    row.guardrail_id = guardrail_id
    row.guardrail_name = name
    row.litellm_params = {"guardrail": provider, "mode": "pre_call"}
    row.guardrail_info = {"type": "ContentSafety", "description": "db-defined"}
    return row


def _metric(guardrail_id: str, date: str = "2026-04-25", requests: int = 10, passed: int = 8, blocked: int = 2) -> Any:
    m = MagicMock()
    m.guardrail_id = guardrail_id
    m.date = date
    m.requests_evaluated = requests
    m.passed_count = passed
    m.blocked_count = blocked
    m.flagged_count = 0
    return m


def _units_row(
    guardrail_id: str,
    date: str = "2026-04-25",
    team_id: str = "",
    api_key: str = "",
    usage_unit: str = "contentPolicyUnits",
    units: int = 1,
) -> Any:
    r = MagicMock()
    r.guardrail_id = guardrail_id
    r.date = date
    r.team_id = team_id
    r.api_key = api_key
    r.usage_unit = usage_unit
    r.units = units
    return r


def _prisma(
    *,
    find_many=None,
    find_unique=None,
    metrics=None,
    index_find_many=None,
    units=None,
) -> MagicMock:
    client = MagicMock()
    db = client.db
    db.litellm_guardrailstable.find_many = AsyncMock(return_value=find_many or [])
    db.litellm_guardrailstable.find_unique = AsyncMock(return_value=find_unique)
    db.litellm_dailyguardrailmetrics.find_many = AsyncMock(return_value=metrics or [])
    db.litellm_dailyguardrailusageunits.find_many = AsyncMock(return_value=units or [])
    db.litellm_spendlogguardrailindex.find_many = AsyncMock(return_value=index_find_many or [])
    db.litellm_spendlogguardrailindex.count = AsyncMock(return_value=0)
    db.litellm_spendlogs.find_many = AsyncMock(return_value=[])
    return client


def _patches(prisma: MagicMock, handler: InMemoryGuardrailHandler):
    return (
        patch("litellm.proxy.proxy_server.prisma_client", prisma),
        patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", handler),
    )


# ---- detail -----------------------------------------------------------------


@pytest.mark.asyncio
async def test_detail_returns_yaml_guardrail_when_db_misses():
    prisma = _prisma(find_unique=None)
    handler = _config_handler(_yaml_guardrail())
    p1, p2 = _patches(prisma, handler)
    with p1, p2:
        resp = await guardrails_usage_detail(
            guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN
        )
    assert resp.guardrail_id == "yaml-1"
    assert resp.guardrail_name == "yaml-pii"
    assert resp.provider == "presidio"  # coerced from the LitellmParams pydantic model
    assert resp.type == "PII"  # from guardrail_info
    assert resp.description == "yaml-defined"


@pytest.mark.asyncio
async def test_detail_404_when_neither_db_nor_config():
    prisma = _prisma(find_unique=None)
    handler = _config_handler()  # empty
    p1, p2 = _patches(prisma, handler)
    with p1, p2, pytest.raises(HTTPException) as exc:
        await guardrails_usage_detail(guardrail_id="ghost", start_date=START, end_date=END, user_api_key_dict=ADMIN)
    assert exc.value.status_code == 404


@pytest.mark.asyncio
async def test_detail_does_not_surface_db_sourced_in_memory_entry():
    """A stale in-memory entry (source=db, gone from DB) must 404, not resurface."""
    prisma = _prisma(find_unique=None)
    handler = InMemoryGuardrailHandler()
    stale = _yaml_guardrail(guardrail_id="stale-1", name="stale")
    handler.IN_MEMORY_GUARDRAILS["stale-1"] = stale
    handler._sources["stale-1"] = "db"
    p1, p2 = _patches(prisma, handler)
    with p1, p2, pytest.raises(HTTPException) as exc:
        await guardrails_usage_detail(guardrail_id="stale-1", start_date=START, end_date=END, user_api_key_dict=ADMIN)
    assert exc.value.status_code == 404


@pytest.mark.asyncio
async def test_detail_db_row_still_resolves():
    prisma = _prisma(find_unique=_db_row(guardrail_id="db-1", provider="aim"))
    handler = _config_handler()
    p1, p2 = _patches(prisma, handler)
    with p1, p2:
        resp = await guardrails_usage_detail(
            guardrail_id="db-1", start_date=START, end_date=END, user_api_key_dict=ADMIN
        )
    assert resp.provider == "aim"
    assert resp.type == "ContentSafety"


# ---- overview ---------------------------------------------------------------


@pytest.mark.asyncio
async def test_overview_includes_yaml_guardrail_with_no_metrics():
    """The core bug: a YAML guardrail with zero metrics must still appear as a row."""
    prisma = _prisma(find_many=[])  # no DB guardrails
    handler = _config_handler(_yaml_guardrail())
    p1, p2 = _patches(prisma, handler)
    with p1, p2:
        resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
    rows = [r for r in resp.rows if r.id == "yaml-1"]
    assert len(rows) == 1
    assert rows[0].name == "yaml-pii"
    assert rows[0].provider == "presidio"
    assert rows[0].type == "PII"
    assert rows[0].requestsEvaluated == 0


@pytest.mark.asyncio
async def test_overview_yaml_metrics_matched_by_logical_name():
    """Daily metrics are keyed by logical name; the YAML row must pick them up."""
    prisma = _prisma(
        find_many=[],
        metrics=[_metric("yaml-pii", requests=10, blocked=2)],  # keyed by name, not uuid
    )
    handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"))
    p1, p2 = _patches(prisma, handler)
    with p1, p2:
        resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
    rows = [r for r in resp.rows if r.id == "yaml-uuid"]
    assert len(rows) == 1
    assert rows[0].requestsEvaluated == 10
    assert rows[0].failRate == 20.0
    # must not also emit an orphan row keyed by the logical name
    assert [r for r in resp.rows if r.id == "yaml-pii"] == []


@pytest.mark.asyncio
async def test_overview_excludes_db_sourced_in_memory_entry():
    """union must not resurrect a stale db-sourced in-memory guardrail."""
    prisma = _prisma(find_many=[])
    handler = InMemoryGuardrailHandler()
    handler.IN_MEMORY_GUARDRAILS["cfg"] = _yaml_guardrail(guardrail_id="cfg", name="cfg-guard")
    handler._sources["cfg"] = "config"
    handler.IN_MEMORY_GUARDRAILS["stale"] = _yaml_guardrail(guardrail_id="stale", name="stale-guard")
    handler._sources["stale"] = "db"
    p1, p2 = _patches(prisma, handler)
    with p1, p2:
        resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
    ids = {r.id for r in resp.rows}
    assert "cfg" in ids
    assert "stale" not in ids


@pytest.mark.asyncio
async def test_overview_reports_usage_units_per_row_and_total():
    """LIT-5650: billable units must surface per guardrail row (matched by
    logical name like the daily metrics) and as a response-level total."""
    prisma = _prisma(
        find_many=[],
        metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)],
        units=[
            _units_row("yaml-pii", usage_unit="topicPolicyUnits", units=4),
            _units_row("yaml-pii", usage_unit="contentPolicyUnits", units=3),
            _units_row("yaml-pii", team_id="team-a", usage_unit="contentPolicyUnits", units=2),
            _units_row("other-guard", usage_unit="topicPolicyUnits", units=7),
        ],
    )
    handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"))
    p1, p2 = _patches(prisma, handler)
    with p1, p2:
        resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
    row = next(r for r in resp.rows if r.id == "yaml-uuid")
    assert row.usageUnits == {"topicPolicyUnits": 4, "contentPolicyUnits": 5}
    assert resp.totalUsageUnits == {"topicPolicyUnits": 11, "contentPolicyUnits": 5}
    units_where = prisma.db.litellm_dailyguardrailusageunits.find_many.call_args.kwargs["where"]
    assert units_where == {"date": {"gte": START, "lte": END}}


@pytest.mark.asyncio
async def test_detail_breaks_units_down_by_day_team_and_key():
    prisma = _prisma(
        find_unique=None,
        units=[
            _units_row("yaml-pii", date="2026-04-25", team_id="team-a", api_key="hash-1", units=2),
            _units_row("yaml-pii", date="2026-04-25", team_id="", api_key="hash-2", units=1),
            _units_row(
                "yaml-pii", date="2026-04-24", team_id="team-a", api_key="hash-1", usage_unit="topicPolicyUnits"
            ),
        ],
    )
    handler = _config_handler(_yaml_guardrail())
    p1, p2 = _patches(prisma, handler)
    with p1, p2:
        resp = await guardrails_usage_detail(
            guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN
        )
    assert resp.usage_units == {"contentPolicyUnits": 3, "topicPolicyUnits": 1}
    assert [p.model_dump() for p in resp.usage_units_daily] == [
        {"date": "2026-04-24", "units": {"topicPolicyUnits": 1}},
        {"date": "2026-04-25", "units": {"contentPolicyUnits": 3}},
    ]
    assert resp.usage_units_by_team == {
        "team-a": {"contentPolicyUnits": 2, "topicPolicyUnits": 1},
        "": {"contentPolicyUnits": 1},
    }
    assert resp.usage_units_by_key == {
        "hash-1": {"contentPolicyUnits": 2, "topicPolicyUnits": 1},
        "hash-2": {"contentPolicyUnits": 1},
    }
    units_where = prisma.db.litellm_dailyguardrailusageunits.find_many.call_args.kwargs["where"]
    assert units_where == {"guardrail_id": {"in": ["yaml-pii", "yaml-1"]}, "date": {"gte": START, "lte": END}}


def _units_table_missing() -> TableNotFoundError:
    return TableNotFoundError(
        data={"user_facing_error": {"meta": {"table": "public.LiteLLM_DailyGuardrailUsageUnits"}}}
    )


@pytest.mark.asyncio
async def test_overview_degrades_units_to_empty_when_units_table_is_missing():
    prisma = _prisma(metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)])
    prisma.db.litellm_dailyguardrailusageunits.find_many = AsyncMock(side_effect=_units_table_missing())
    handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"))
    p1, p2 = _patches(prisma, handler)
    with p1, p2:
        resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
    row = next(r for r in resp.rows if r.id == "yaml-uuid")
    assert (row.requestsEvaluated, row.usageUnits) == (4, {})
    assert (resp.totalRequests, resp.totalBlocked, resp.totalUsageUnits) == (4, 1, {})


@pytest.mark.asyncio
async def test_detail_degrades_units_to_empty_when_units_table_is_missing():
    prisma = _prisma(metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)])
    prisma.db.litellm_dailyguardrailusageunits.find_many = AsyncMock(side_effect=_units_table_missing())
    handler = _config_handler(_yaml_guardrail())
    p1, p2 = _patches(prisma, handler)
    with p1, p2:
        resp = await guardrails_usage_detail(
            guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN
        )
    assert (resp.requestsEvaluated, resp.failRate) == (4, 25.0)
    assert (resp.usage_units, list(resp.usage_units_daily), resp.usage_units_by_team, resp.usage_units_by_key) == (
        {},
        [],
        {},
        {},
    )


# ---- logs -------------------------------------------------------------------


@pytest.mark.asyncio
async def test_logs_resolves_config_guardrail_logical_name():
    """The index query must include the YAML guardrail's logical name alias."""
    prisma = _prisma(find_unique=None)
    handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"))
    p1, p2 = _patches(prisma, handler)
    with p1, p2:
        await guardrails_usage_logs(
            guardrail_id="yaml-uuid",
            policy_id=None,
            page=1,
            page_size=50,
            action=None,
            start_date=START,
            end_date=END,
            user_api_key_dict=ADMIN,
        )
    where = prisma.db.litellm_spendlogguardrailindex.find_many.call_args.kwargs["where"]
    assert where["guardrail_id"] == {"in": ["yaml-uuid", "yaml-pii"]}


# ---- date window cap (LIT-5762) ---------------------------------------------


@pytest.mark.asyncio
async def test_overview_rejects_range_over_max_days():
    prisma = _prisma()
    handler = _config_handler()
    p1, p2 = _patches(prisma, handler)
    with p1, p2, pytest.raises(HTTPException) as exc:
        await guardrails_usage_overview(start_date="2020-01-01", end_date=END, user_api_key_dict=ADMIN)
    assert exc.value.status_code == 400
    assert "366" in str(exc.value.detail)


@pytest.mark.asyncio
async def test_overview_accepts_range_at_exactly_max_days():
    prisma = _prisma()
    handler = _config_handler()
    p1, p2 = _patches(prisma, handler)
    with p1, p2:
        resp = await guardrails_usage_overview(start_date="2025-04-26", end_date="2026-04-27", user_api_key_dict=ADMIN)
    assert resp.totalRequests == 0


@pytest.mark.asyncio
async def test_overview_rejects_malformed_dates():
    prisma = _prisma()
    handler = _config_handler()
    p1, p2 = _patches(prisma, handler)
    with p1, p2, pytest.raises(HTTPException) as exc:
        await guardrails_usage_overview(start_date="not-a-date", end_date=END, user_api_key_dict=ADMIN)
    assert exc.value.status_code == 400


@pytest.mark.asyncio
async def test_overview_rejects_non_canonical_date_format():
    prisma = _prisma()
    handler = _config_handler()
    p1, p2 = _patches(prisma, handler)
    with p1, p2, pytest.raises(HTTPException) as exc:
        await guardrails_usage_overview(start_date="20260420", end_date=END, user_api_key_dict=ADMIN)
    assert exc.value.status_code == 400
    assert "YYYY-MM-DD" in str(exc.value.detail)


@pytest.mark.asyncio
async def test_detail_rejects_reversed_dates():
    prisma = _prisma(find_unique=_db_row())
    handler = _config_handler()
    p1, p2 = _patches(prisma, handler)
    with p1, p2, pytest.raises(HTTPException) as exc:
        await guardrails_usage_detail(guardrail_id="db-1", start_date=END, end_date=START, user_api_key_dict=ADMIN)
    assert exc.value.status_code == 400


@pytest.mark.asyncio
async def test_policies_overview_rejects_range_over_max_days():
    prisma = _prisma()
    handler = _config_handler()
    p1, p2 = _patches(prisma, handler)
    with p1, p2, pytest.raises(HTTPException) as exc:
        await policies_usage_overview(start_date="2020-01-01", end_date=END, user_api_key_dict=ADMIN)
    assert exc.value.status_code == 400


@pytest.mark.asyncio
async def test_detail_prev_trend_query_is_bounded():
    """Regression: the trend query scanned every metrics row before start_date."""
    prisma = _prisma(find_unique=_db_row())
    handler = _config_handler()
    p1, p2 = _patches(prisma, handler)
    with p1, p2:
        await guardrails_usage_detail(guardrail_id="db-1", start_date=START, end_date=END, user_api_key_dict=ADMIN)
    wheres = [c.kwargs["where"] for c in prisma.db.litellm_dailyguardrailmetrics.find_many.await_args_list]
    prev_wheres = [w for w in wheres if "lt" in w.get("date", {})]
    assert prev_wheres
    assert all("gte" in w["date"] for w in prev_wheres)