litellm/tests/integration/management/test_key_reset_spend.py

import uuid

import pytest

from litellm.proxy.utils import hash_token

from tests.integration._support.actors import TEAM_ALPHA, TEAM_BETA, Actor
from tests.integration._support.actors import create_scratch_key

pytestmark = pytest.mark.asyncio(loop_scope="session")

_SEED_SPEND = 5.0
_RESET_TO = 2.0


# POST /key/{key}/reset_spend. The target key is pre-seeded with spend=5.0 so
# reset_to=2.0 always clears _validate_reset_spend_value (which runs before
# authz). _check_proxy_or_team_admin_for_key then allows only PROXY_ADMIN or a
# team admin of the key's team — there is no org-admin branch, and a teamless
# "self" key has no team to admin. ORG_ADMIN-role callers are stopped 401 at
# the management-route gate before the handler runs.
_SCENARIOS = [
    ("self/proxy_admin", Actor.PROXY_ADMIN, "self", 200),
    ("self/org_admin", Actor.ORG_ADMIN, "self", 401),
    ("self/team_admin", Actor.TEAM_ADMIN, "self", 403),
    ("self/internal_user", Actor.INTERNAL_USER, "self", 403),
    ("self/cross_org_user", Actor.CROSS_ORG_USER, "self", 403),
    ("team_alpha/proxy_admin", Actor.PROXY_ADMIN, "team_alpha", 200),
    ("team_alpha/org_admin", Actor.ORG_ADMIN, "team_alpha", 401),
    ("team_alpha/team_admin", Actor.TEAM_ADMIN, "team_alpha", 200),
    ("team_alpha/internal_user", Actor.INTERNAL_USER, "team_alpha", 403),
    ("team_alpha/owner", Actor.OWNER, "team_alpha", 403),
    ("team_alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "team_alpha", 403),
    ("team_alpha/cross_org_user", Actor.CROSS_ORG_USER, "team_alpha", 403),
    ("team_alpha/service_account", Actor.SERVICE_ACCOUNT, "team_alpha", 403),
    ("team_alpha/org_b_admin", Actor.ORG_B_ADMIN, "team_alpha", 401),
    ("team_beta/proxy_admin", Actor.PROXY_ADMIN, "team_beta", 200),
    ("team_beta/org_admin", Actor.ORG_ADMIN, "team_beta", 401),
    ("team_beta/team_admin", Actor.TEAM_ADMIN, "team_beta", 403),
    ("team_beta/cross_org_user", Actor.CROSS_ORG_USER, "team_beta", 403),
    ("team_beta/org_b_admin", Actor.ORG_B_ADMIN, "team_beta", 401),
]


async def _seed_target(gateway_client, seeder, prefix, actors, shape, caller) -> str:
    if shape == "self":
        return await create_scratch_key(
            gateway_client, seeder, prefix, user_id=caller.user_id
        )
    if shape == "team_alpha":
        return await create_scratch_key(
            gateway_client,
            seeder,
            prefix,
            user_id=actors.keys[Actor.OWNER].user_id,
            team_id=actors.team_alpha_id,
        )
    if shape == "team_beta":
        return await create_scratch_key(
            gateway_client,
            seeder,
            prefix,
            user_id=actors.keys[Actor.CROSS_ORG_USER].user_id,
            team_id=actors.team_beta_id,
        )
    pytest.fail(f"unknown shape={shape}")  # pragma: no cover


@pytest.mark.parametrize(
    "actor,shape,expected_status",
    [(a, sh, s) for (_id, a, sh, s) in _SCENARIOS],
    ids=[s[0] for s in _SCENARIOS],
)
async def test_key_reset_spend_authz_matrix(
    actor: Actor,
    shape: str,
    expected_status: int,
    gateway_client,
    db,
    resources,
    actors,
):
    caller = actors.keys[actor]
    seeder = actors.keys[Actor.PROXY_ADMIN].cleartext
    target = await _seed_target(
        gateway_client, seeder, resources.prefix, actors, shape, caller
    )
    hashed = hash_token(target)
    await db.db.litellm_verificationtoken.update(
        where={"token": hashed}, data={"spend": _SEED_SPEND}
    )

    resp = await gateway_client.post(
        f"/key/{target}/reset_spend",
        headers={"Authorization": f"Bearer {caller.cleartext}"},
        json={"reset_to": _RESET_TO},
    )
    assert (
        resp.status_code == expected_status
    ), f"{actor.value} {shape}: {resp.status_code} {resp.text}"

    row = await db.db.litellm_verificationtoken.find_unique(where={"token": hashed})
    assert row is not None
    if expected_status == 200:
        assert row.spend == _RESET_TO
    else:
        assert row.spend == _SEED_SPEND, "denied but spend reset"


@pytest.mark.parametrize(
    "actor", [Actor.PROXY_ADMIN, Actor.TEAM_ADMIN], ids=["proxy_admin", "team_admin"]
)
async def test_key_reset_spend_missing_key_is_404(actor: Actor, gateway_client, actors):
    """A well-formed but unseeded key is 404 before any spend validation."""
    resp = await gateway_client.post(
        f"/key/sk-{uuid.uuid4().hex}/reset_spend",
        headers={"Authorization": f"Bearer {actors.keys[actor].cleartext}"},
        json={"reset_to": 0.0},
    )
    assert resp.status_code == 404, resp.text


async def test_key_reset_spend_above_current_spend_is_400(
    gateway_client, db, resources, actors
):
    """reset_to above the key's current spend is rejected 400."""
    admin = actors.keys[Actor.PROXY_ADMIN]
    target = await create_scratch_key(
        gateway_client, admin.cleartext, resources.prefix, user_id=admin.user_id
    )
    resp = await gateway_client.post(
        f"/key/{target}/reset_spend",
        headers={"Authorization": f"Bearer {admin.cleartext}"},
        json={"reset_to": 1.0},
    )
    assert resp.status_code == 400, resp.text