litellm/tests/litellm_utils_tests/test_proxy_budget_reset.py

import asyncio
from datetime import datetime, timedelta, timezone
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from dotenv import load_dotenv

load_dotenv()

from litellm.proxy._types import LiteLLM_BudgetTableFull


from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob

# Note: In our "fake" items we use dicts with fields that our fake reset functions modify.
# In a real-world scenario, these would be instances of LiteLLM_VerificationToken, LiteLLM_UserTable, etc.


def _attrify(d: dict):
    """
    Wrap a dict so that attribute access (`.token`, `.user_id`, `.team_id`,
    etc.) works alongside the existing item-access the fake_reset_* helpers
    rely on. The reset job's narrow-write helpers use `getattr(item, "token",
    None)` (et al), which returns None for plain dicts — that would silently
    skip the row.
    """

    class _AttrDict(dict):
        def __getattr__(self, k):
            try:
                return self[k]
            except KeyError:
                raise AttributeError(k)

        def __setattr__(self, k, v):
            self[k] = v

    return _AttrDict(d)


def _wire_batcher_for_test(prisma_client, fail_commit=False):
    """
    Wire prisma_client.db.batch_() to return a mock batcher whose .commit() is
    awaitable and whose per-table .update()/.update_many() calls get captured.
    The reset job writes every reset through prisma.db.batch_() — key/user/team
    rows one by one, and the budget tier's cascade as a single transaction — so
    tests must let that batch path complete.

    Only committed batches contribute to the returned list, mirroring prisma:
    with fail_commit=True the transaction blows up and must persist nothing.

    Returns the list that will accumulate {table, op, where, data} dicts from
    each captured write.
    """
    batch_calls = []

    def make_batcher():
        queued = []

        class _Table:
            def __init__(self, table_name):
                self._table_name = table_name

            def update(self, where=None, data=None):
                queued.append(
                    {
                        "table": self._table_name,
                        "op": "update",
                        "where": where,
                        "data": data,
                    }
                )

            def update_many(self, where=None, data=None):
                queued.append(
                    {
                        "table": self._table_name,
                        "op": "update_many",
                        "where": where,
                        "data": data,
                    }
                )

        async def commit():
            if fail_commit:
                raise RuntimeError("simulated Postgres failure committing the batch")
            batch_calls.extend(queued)

        batcher = MagicMock()
        batcher.litellm_verificationtoken = _Table("key")
        batcher.litellm_usertable = _Table("user")
        batcher.litellm_teamtable = _Table("team")
        batcher.litellm_budgettable = _Table("budget")
        batcher.litellm_teammembership = _Table("team_membership")
        batcher.litellm_organizationtable = _Table("org")
        batcher.litellm_tagtable = _Table("tag")
        batcher.litellm_endusertable = _Table("enduser")
        batcher.commit = commit
        return batcher

    prisma_client.db.batch_ = MagicMock(side_effect=make_batcher)
    return batch_calls


def _wire_cascade_reads_for_test(prisma_client, endusers=()):
    """
    The budget tier's cascade reads the rows it is about to zero, so their
    spend counters can be invalidated after the commit. Give each of those
    tables an awaitable find_many so the reads resolve instead of falling into
    the job's warn-and-continue path.

    End users are read by the post-commit invalidation walk rather than by
    ``get_data``, so callers that care about customers pass them here.
    """
    for table in (
        "litellm_teammembership",
        "litellm_verificationtoken",
        "litellm_organizationtable",
        "litellm_tagtable",
    ):
        getattr(prisma_client.db, table).find_many = AsyncMock(return_value=[])
    prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=list(endusers))








@pytest.mark.asyncio
async def test_reset_budget_endusers_are_zeroed_with_the_budget_window_advance():
    """
    The happy path: every end user the tier gates is zeroed and the tier's
    budget_reset_at advances, all inside one transaction.
    """
    endusers = [
        _attrify({"user_id": f"user{i}", "spend": 20.0 + i, "budget_id": "budget1"})
        for i in range(1, 7)
    ]

    budget1 = LiteLLM_BudgetTableFull(
        **{
            "budget_id": "budget1",
            "max_budget": 65.0,
            "budget_duration": "2d",
            "created_at": datetime.now(timezone.utc) - timedelta(days=3),
        }
    )

    prisma_client = MagicMock()

    async def get_data_mock(table_name, *args, **kwargs):
        if table_name == "budget":
            return [budget1]
        elif table_name == "enduser":
            return endusers
        return []

    prisma_client.get_data = AsyncMock()
    prisma_client.get_data.side_effect = get_data_mock
    prisma_client.update_data = AsyncMock()
    batch_calls = _wire_batcher_for_test(prisma_client)

    proxy_logging_obj = MagicMock()
    proxy_logging_obj.service_logging_obj = MagicMock()
    proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
    proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()

    job = ResetBudgetJob(proxy_logging_obj, prisma_client)

    await job.reset_budget_for_litellm_budget_table()
    await asyncio.sleep(0.1)

    assert prisma_client.db.batch_.call_count == 1, "the cascade must be one transaction"

    enduser_writes = [c for c in batch_calls if c["table"] == "enduser"]
    assert len(enduser_writes) == 1
    assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}}
    assert enduser_writes[0]["data"] == {"spend": 0}

    budget_writes = [c for c in batch_calls if c["table"] == "budget"]
    assert len(budget_writes) == 1
    assert budget_writes[0]["where"] == {"budget_id": "budget1"}
    assert budget_writes[0]["data"]["budget_reset_at"] > datetime.now(timezone.utc)

    proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called()




@pytest.mark.asyncio
async def test_reset_budget_continues_other_categories_on_failure():
    """
    Test that executing the overall reset_budget() method continues to process keys, users, and teams,
    even if one of the sub-categories (here, users) experiences a partial failure.

    In this simulation:
      - All keys are processed successfully.
      - One of the two users fails.
      - All teams are processed successfully.

    We then assert that:
      - update_data is called for each category with the correctly updated items.
      - Each get_data call is made (indicating that one failing category did not abort the others).
    """
    # Arrange dummy items for each table
    key1 = {"id": "key1", "spend": 10.0, "budget_duration": 60}
    key2 = {"id": "key2", "spend": 15.0, "budget_duration": 60}
    user1 = {
        "id": "user1",
        "spend": 20.0,
        "budget_duration": 120,
    }  # Will fail in user reset
    user2 = {"id": "user2", "spend": 25.0, "budget_duration": 120}  # Succeeds
    team1 = {"id": "team1", "spend": 30.0, "budget_duration": 180}
    team2 = {"id": "team2", "spend": 35.0, "budget_duration": 180}
    enduser1 = {"user_id": "user1", "spend": 25.0, "budget_id": "budget1"}
    budget1 = LiteLLM_BudgetTableFull(
        **{
            "budget_id": "budget1",
            "max_budget": 65.0,
            "budget_duration": "2d",
            "created_at": datetime.now(timezone.utc) - timedelta(days=3),
        }
    )

    prisma_client = MagicMock()

    async def fake_get_data(*, table_name, query_type, **kwargs):
        if table_name == "key":
            return [key1, key2]
        elif table_name == "user":
            return [user1, user2]
        elif table_name == "team":
            return [team1, team2]
        elif table_name == "budget":
            return [budget1]
        elif table_name == "enduser":
            return [enduser1]
        return []

    prisma_client.get_data = AsyncMock(side_effect=fake_get_data)
    prisma_client.update_data = AsyncMock()
    batch_calls = _wire_batcher_for_test(prisma_client)
    # ID fields required by the new write path's where clauses; _AttrDict
    # lets getattr() see them alongside the item-access fake_reset_* helpers use.
    for k in [key1, key2]:
        k.setdefault("token", k["id"])
    for u in [user1, user2]:
        u.setdefault("user_id", u["id"])
    for t in [team1, team2]:
        t.setdefault("team_id", t["id"])
    key1, key2 = _attrify(key1), _attrify(key2)
    user1, user2 = _attrify(user1), _attrify(user2)
    team1, team2 = _attrify(team1), _attrify(team2)
    enduser1 = _attrify(enduser1)
    pre_reset_spend = {
        **{k["token"]: k["spend"] for k in [key1, key2]},
        **{u["user_id"]: u["spend"] for u in [user2]},
        **{t["team_id"]: t["spend"] for t in [team1, team2]},
    }
    _wire_cascade_reads_for_test(prisma_client, endusers=[enduser1])

    proxy_logging_obj = MagicMock()
    proxy_logging_obj.service_logging_obj = MagicMock()
    proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
    proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()

    job = ResetBudgetJob(proxy_logging_obj, prisma_client)

    async def fake_reset_key(key, current_time, reset_settings=None):
        key["spend"] = 0.0
        key["budget_reset_at"] = (
            current_time + timedelta(seconds=key["budget_duration"])
        ).isoformat()
        return key

    async def fake_reset_user(user, current_time, reset_settings=None):
        if user["id"] == "user1":
            raise Exception("Simulated failure for user1")
        user["spend"] = 0.0
        user["budget_reset_at"] = (
            current_time + timedelta(seconds=user["budget_duration"])
        ).isoformat()
        return user

    async def fake_reset_team(team, current_time, reset_settings=None):
        team["spend"] = 0.0
        team["budget_reset_at"] = (
            current_time + timedelta(seconds=team["budget_duration"])
        ).isoformat()
        return team

    with (
        patch.object(
            ResetBudgetJob, "_reset_budget_for_key", side_effect=fake_reset_key
        ) as mock_reset_key,
        patch.object(
            ResetBudgetJob, "_reset_budget_for_user", side_effect=fake_reset_user
        ) as mock_reset_user,
        patch.object(
            ResetBudgetJob, "_reset_budget_for_team", side_effect=fake_reset_team
        ) as mock_reset_team,
    ):
        # Call the overall reset_budget method.
        await job.reset_budget()
        await asyncio.sleep(0.1)

    # Verify that get_data was called for each table. We can check the table names across calls.
    called_tables = {
        call.kwargs.get("table_name") for call in prisma_client.get_data.await_args_list
    }
    assert called_tables == {"key", "user", "team", "budget"}
    # Customers are not part of that set: the cascade zeroes them by budget link
    # and reads them only afterwards, to invalidate their cached spend.
    prisma_client.db.litellm_endusertable.find_many.assert_awaited()

    # Every category writes through the batch path now, so update_data is unused.
    prisma_client.update_data.assert_not_awaited()

    # The budget tier's cascade still ran despite the failing user category.
    assert len([c for c in batch_calls if c["table"] == "team_membership"]) == 1
    enduser_writes = [c for c in batch_calls if c["table"] == "enduser"]
    assert len(enduser_writes) == 1
    assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}}
    assert enduser_writes[0]["data"] == {"spend": 0}

    # Check the new batch write path: 2 keys + 1 user (user1 failed) + 2 teams.
    # `op` separates the per-row resets from the cascade sweep, which also
    # targets the key table.
    key_writes = [c for c in batch_calls if c["table"] == "key" and c["op"] == "update"]
    user_writes = [c for c in batch_calls if c["table"] == "user"]
    team_writes = [c for c in batch_calls if c["table"] == "team"]
    assert len(key_writes) == 2
    assert len(user_writes) == 1
    assert user_writes[0]["where"] == {"user_id": "user2"}
    assert len(team_writes) == 2
    # Every batched write must carry only the two reset fields, never the full row.
    for c in key_writes + user_writes + team_writes:
        assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
        assert c["data"]["spend"] == {
            "decrement": pre_reset_spend[next(iter(c["where"].values()))]
        }


# ---------------------------------------------------------------------------
# Additional tests for service logger behavior (keys, users, teams, endusers)
# ---------------------------------------------------------------------------