litellm/tests/unit/proxy/db/test_model_usage_rollup.py

import asyncio
from datetime import datetime, timezone
from typing import Any
from unittest.mock import MagicMock

import httpx
import pytest

from litellm.proxy.db import rollup_lock_timeout
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.db.model_usage_rollup import (
    ModelUsageKey,
    ModelUsageTransaction,
    build_model_usage_transaction,
    flush_model_usage_transactions,
    model_usage_task_type,
)
from litellm.proxy.db.rollup_lock_timeout import ROLLUP_LOCK_TIMEOUT_SQL


class _FakeBatcher:
    def __init__(self) -> None:
        self.litellm_dailymodelusage = MagicMock()
        self.execute_raw = MagicMock()

    async def __aenter__(self) -> "_FakeBatcher":
        return self

    async def __aexit__(self, *args: Any) -> None:
        return None


def _prisma(batch_: MagicMock) -> MagicMock:
    prisma = MagicMock()
    prisma.db.batch_ = batch_
    return prisma


def _payload(**overrides: Any) -> dict[str, Any]:
    return {
        "spend": 0.25,
        "prompt_tokens": 10,
        "completion_tokens": 20,
        "startTime": datetime(2026, 9, 28, 13, tzinfo=timezone.utc),
        "model": "openai/gpt-5.4-mini",
        "model_group": "fast-chat",
        "metadata": "{}",
        "request_tags": "[]",
        "custom_llm_provider": "openai",
        "status": "success",
        **overrides,
    }


def _transaction(model: str, spend: float, successful: bool = True) -> ModelUsageTransaction:
    return ModelUsageTransaction(
        key=ModelUsageKey(
            date="2026-09-28", model_group=model, model=model, custom_llm_provider="openai", task_type="debugging"
        ),
        spend=spend,
        prompt_tokens=10,
        completion_tokens=5,
        successful=successful,
    )


async def _no_sleep(seconds: float) -> None:
    return None


def test_model_usage_task_type_reads_task_tag_or_defaults() -> None:
    assert model_usage_task_type('["team-a", "task:classification"]') == "classification"
    assert model_usage_task_type('["task:made-up"]') == "uncategorized"
    assert model_usage_task_type('["debugging"]') == "uncategorized"
    assert model_usage_task_type("[]") == "uncategorized"
    assert model_usage_task_type("not json") == "uncategorized"


def test_build_model_usage_transaction_keys_on_day_model_and_task() -> None:
    transaction = build_model_usage_transaction(_payload(request_tags='["task:debugging"]', status="failure"))

    assert transaction == ModelUsageTransaction(
        key=ModelUsageKey(
            date="2026-09-28",
            model_group="fast-chat",
            model="openai/gpt-5.4-mini",
            custom_llm_provider="openai",
            task_type="debugging",
        ),
        spend=0.25,
        prompt_tokens=10,
        completion_tokens=20,
        successful=False,
    )


def test_build_model_usage_transaction_falls_back_for_missing_model_fields() -> None:
    transaction = build_model_usage_transaction(
        _payload(model="", model_group=None, custom_llm_provider=None, startTime="2026-09-28T01:02:03Z")
    )

    assert transaction is not None
    assert transaction.key == ModelUsageKey(
        date="2026-09-28",
        model_group="unknown",
        model="unknown",
        custom_llm_provider="unknown",
        task_type="uncategorized",
    )


@pytest.mark.parametrize(
    "overrides",
    [{"metadata": '{"internal_call_origin": "health_check"}'}, {"startTime": "bad"}],
)
def test_build_model_usage_transaction_skips_internal_calls_and_bad_dates(overrides: dict[str, Any]) -> None:
    assert build_model_usage_transaction(_payload(**overrides)) is None


@pytest.mark.asyncio
async def test_flush_aggregates_each_rollup_row_into_one_upsert() -> None:
    batcher = _FakeBatcher()
    prisma = _prisma(MagicMock(return_value=batcher))

    await flush_model_usage_transactions(
        prisma_client=prisma,
        transactions=[
            _transaction("gpt-5", 0.5),
            _transaction("claude", 1.0),
            _transaction("gpt-5", 0.25, successful=False),
            _transaction("gpt-5", 0.25),
        ],
    )

    upserts = {
        call.kwargs["where"]["date_model_group_model_custom_llm_provider_task_type"]["model"]: call.kwargs["data"]
        for call in batcher.litellm_dailymodelusage.upsert.call_args_list
    }
    assert list(upserts) == ["claude", "gpt-5"]
    gpt = upserts["gpt-5"]
    assert gpt["create"]["spend"] == 1.0
    assert gpt["create"]["prompt_tokens"] == 30
    assert gpt["create"]["completion_tokens"] == 15
    assert gpt["create"]["request_count"] == 3
    assert gpt["create"]["successful_requests"] == 2
    assert gpt["create"]["failed_requests"] == 1
    assert gpt["update"] == {
        "spend": {"increment": 1.0},
        "prompt_tokens": {"increment": 30},
        "completion_tokens": {"increment": 15},
        "request_count": {"increment": 3},
        "successful_requests": {"increment": 2},
        "failed_requests": {"increment": 1},
    }
    assert upserts["claude"]["create"]["request_count"] == 1


@pytest.mark.asyncio
async def test_flush_with_no_transactions_touches_nothing() -> None:
    prisma = _prisma(MagicMock())
    await flush_model_usage_transactions(prisma_client=prisma, transactions=[])
    prisma.db.batch_.assert_not_called()


@pytest.mark.asyncio
async def test_flush_retries_connection_errors(monkeypatch: pytest.MonkeyPatch) -> None:
    batcher = _FakeBatcher()
    prisma = _prisma(MagicMock(side_effect=[httpx.ConnectError("down"), batcher]))
    monkeypatch.setattr("litellm.proxy.db.model_usage_rollup.asyncio.sleep", _no_sleep)

    await flush_model_usage_transactions(prisma_client=prisma, transactions=[_transaction("gpt-5", 0.1)])

    assert prisma.db.batch_.call_count == 2
    batcher.litellm_dailymodelusage.upsert.assert_called_once()


@pytest.mark.asyncio
async def test_flush_does_not_retry_ambiguous_errors() -> None:
    prisma = _prisma(MagicMock(side_effect=httpx.ReadTimeout("ambiguous")))

    with pytest.raises(httpx.ReadTimeout):
        await flush_model_usage_transactions(prisma_client=prisma, transactions=[_transaction("gpt-5", 0.1)])

    prisma.db.batch_.assert_called_once()


@pytest.mark.asyncio
async def test_request_time_path_queues_usage_instead_of_writing_to_the_db() -> None:
    prisma = MagicMock()
    prisma.model_usage_transactions = []
    prisma._model_usage_transactions_lock = asyncio.Lock()

    await DBSpendUpdateWriter()._batch_database_updates(
        response_cost=0.25,
        user_id="u1",
        hashed_token="t1",
        team_id=None,
        org_id=None,
        end_user_id=None,
        prisma_client=prisma,
        litellm_proxy_budget_name=None,
        payload=_payload(request_id="req-1"),
    )

    assert [transaction.key.model for transaction in prisma.model_usage_transactions] == ["openai/gpt-5.4-mini"]
    prisma.db.litellm_dailymodelusage.upsert.assert_not_called()


@pytest.mark.asyncio
async def test_rollup_transaction_sets_the_lock_timeout_before_its_first_upsert(monkeypatch) -> None:
    """The DailyModelUsage rollup shares the DailyToolSpend write shape, so it runs
    under the same per-transaction lock budget, issued before any upsert queues."""
    monkeypatch.setattr(rollup_lock_timeout, "SPEND_ROLLUP_LOCK_TIMEOUT_MS", 250)
    batcher = _FakeBatcher()
    prisma = _prisma(MagicMock(return_value=batcher))
    order: list[str] = []
    batcher.execute_raw.side_effect = lambda *_: order.append("lock_timeout")
    batcher.litellm_dailymodelusage.upsert.side_effect = lambda **_: order.append("upsert")

    await flush_model_usage_transactions(
        prisma_client=prisma,
        transactions=[
            build_model_usage_transaction(_payload()),
            build_model_usage_transaction(_payload(model="other")),
        ],
    )

    assert batcher.execute_raw.call_args_list == [((ROLLUP_LOCK_TIMEOUT_SQL, "250ms"),)]
    assert order == ["lock_timeout", "upsert", "upsert"]