litellm/tests/unit/proxy/auth/test_auth_object_prefetch.py

"""Counts the Redis round trips and DB queries auth object reads cost per cache regime, and checks that the
per-object getters still enforce on their own when the prefetch cannot help."""

from unittest.mock import AsyncMock, MagicMock

import pytest
from fastapi import HTTPException

from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy._types import (
    LiteLLM_OrganizationTable,
    LiteLLM_TeamMembership,
    LiteLLM_TeamTableCachedObj,
    LiteLLM_UserTable,
    UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_checks import (
    get_end_user_object,
    get_org_object,
    get_team_membership,
    get_team_object,
    get_user_object,
)
from litellm.proxy.auth.auth_object_prefetch import AuthObjectRefs, prefetch_auth_objects, prefetch_identity_keys
from litellm.proxy.common_utils.user_api_key_cache import (
    UserApiKeyCache,
    end_user_cache_key,
    end_user_restricted_registry_cache_key,
)
from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache

USER_ID = "prefetch-user"
TEAM_ID = "prefetch-team"
ORG_ID = "prefetch-org"

USER_ROW = {
    "user_id": USER_ID,
    "max_budget": 50.0,
    "spend": 1.0,
    "models": ["gpt-5.4-mini"],
    "organization_memberships": [],
}
TEAM_ROW = {
    "team_id": TEAM_ID,
    "organization_id": ORG_ID,
    "max_budget": 500.0,
    "spend": 2.0,
    "models": [],
    "blocked": False,
    "members_with_roles": {},
}
MEMBERSHIP_ROW = {
    "user_id": USER_ID,
    "team_id": TEAM_ID,
    "spend": 3.0,
    "budget_id": "b1",
    "litellm_budget_table": {"budget_id": "b1", "max_budget": 20.0},
}
ORG_ROW = {
    "organization_id": ORG_ID,
    "organization_alias": "org",
    "budget_id": "b2",
    "created_by": "admin",
    "updated_by": "admin",
    "models": [],
    "spend": 4.0,
    "litellm_budget_table": {"budget_id": "b2", "max_budget": 1000.0},
}
ALL_ROWS = {
    "user_row": USER_ROW,
    "team_row": TEAM_ROW,
    "membership_row": MEMBERSHIP_ROW,
    "organization_row": ORG_ROW,
    "project_row": None,
}


def _prisma(rows: dict[str, object] | None = ALL_ROWS) -> MagicMock:
    prisma = MagicMock(name="prisma_client")
    prisma.db.query_first = AsyncMock(return_value=rows)
    return prisma


def _non_prefetch_db_calls(prisma: MagicMock) -> list[str]:
    return [str(call) for call in prisma.db.mock_calls if not str(call).startswith("call.query_first(")]


def _cache(redis: FakeAuthCache | None) -> UserApiKeyCache:
    """A worker of its own (empty memory) on ``redis``'s Redis, or a worker without Redis."""
    if redis is None:
        return UserApiKeyCache(in_memory_cache=InMemoryCache())
    return shared_cache(FakeAuthCache(redis=redis.redis, commands=redis.commands, fail=redis.fail))[0]


def _round_trips(redis: FakeAuthCache) -> int:
    return len(redis.commands)


def _writes(redis: FakeAuthCache) -> list[str]:
    return sorted(
        f"SET {key} ttl={ttl}" for command in redis.commands if command[0] == "WRITE_BACK" for key, ttl in command[1]
    )


def _refs() -> AuthObjectRefs:
    return AuthObjectRefs.from_token(UserAPIKeyAuth(token="t", user_id=USER_ID, team_id=TEAM_ID, org_id=ORG_ID))


async def _read_all_through_getters(
    cache: UserApiKeyCache, prisma: MagicMock
) -> tuple[
    LiteLLM_UserTable | None,
    LiteLLM_TeamTableCachedObj,
    LiteLLM_TeamMembership | None,
    LiteLLM_OrganizationTable | None,
]:
    return (
        await get_user_object(user_id=USER_ID, prisma_client=prisma, user_api_key_cache=cache, user_id_upsert=False),
        await get_team_object(team_id=TEAM_ID, prisma_client=prisma, user_api_key_cache=cache),
        await get_team_membership(user_id=USER_ID, team_id=TEAM_ID, prisma_client=prisma, user_api_key_cache=cache),
        await get_org_object(org_id=ORG_ID, prisma_client=prisma, user_api_key_cache=cache, include_budget_table=True),
    )


def test_refs_from_token_only_names_membership_when_both_ids_present():
    assert AuthObjectRefs.from_token(UserAPIKeyAuth(token="t", team_id=TEAM_ID)).membership_user_id is None
    assert AuthObjectRefs.from_token(UserAPIKeyAuth(token="t", user_id=USER_ID)).membership_user_id is None
    assert AuthObjectRefs.from_token(UserAPIKeyAuth(token="t", user_id=USER_ID, team_id=TEAM_ID)) == AuthObjectRefs(
        user_id=USER_ID, team_id=TEAM_ID, membership_user_id=USER_ID
    )


@pytest.mark.asyncio
async def test_cold_regime_is_one_mget_one_query_and_the_getters_never_touch_io_again():
    redis = FakeAuthCache()
    prisma = _prisma()
    cache = _cache(redis)

    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma)

    assert prisma.db.query_first.await_count == 1
    assert prisma.db.query_first.await_args.args[1:] == (USER_ID, TEAM_ID, USER_ID, ORG_ID, None)
    mgets = [c for c in redis.commands if c[0] == "MGET"]
    assert len(mgets) == 1
    assert set(mgets[0][1]) == {
        USER_ID,
        f"team_id:{TEAM_ID}",
        f"{TEAM_ID}_{USER_ID}",
        f"team_membership:{USER_ID}:{TEAM_ID}",
        f"org_id:{ORG_ID}",
        f"org_id:{ORG_ID}:with_budget",
    }
    assert _writes(redis) == sorted(
        [
            f"SET {TEAM_ID}_{USER_ID} ttl=5",
            f"SET org_id:{ORG_ID} ttl=5",
            f"SET org_id:{ORG_ID}:with_budget ttl=5",
            f"SET {USER_ID} ttl=60",
            f"SET team_id:{TEAM_ID} ttl=60",
            f"SET team_membership:{USER_ID}:{TEAM_ID} ttl=None",
        ]
    )
    assert _round_trips(redis) == 2, "one MGET, one pipeline"

    before = (_round_trips(redis), prisma.db.query_first.await_count)
    user, team, membership, org = await _read_all_through_getters(cache, prisma)
    assert (_round_trips(redis), prisma.db.query_first.await_count) == before
    assert _non_prefetch_db_calls(prisma) == []

    assert isinstance(user, LiteLLM_UserTable) and user.max_budget == 50.0
    assert isinstance(team, LiteLLM_TeamTableCachedObj) and team.organization_id == ORG_ID
    assert team.last_refreshed_at is not None
    assert isinstance(membership, LiteLLM_TeamMembership) and membership.litellm_budget_table is not None
    assert membership.litellm_budget_table.max_budget == 20.0
    assert isinstance(org, LiteLLM_OrganizationTable) and org.litellm_budget_table is not None
    assert org.litellm_budget_table.max_budget == 1000.0


@pytest.mark.asyncio
async def test_redis_warm_regime_is_exactly_one_mget_and_zero_queries():
    seeded = FakeAuthCache()
    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=_prisma())

    redis = FakeAuthCache(redis=seeded.redis)
    prisma = _prisma()
    cache = _cache(redis)
    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma)

    assert _round_trips(redis) == 1
    assert redis.commands[0][0] == "MGET"
    assert prisma.db.query_first.await_count == 0

    user, team, membership, org = await _read_all_through_getters(cache, prisma)
    assert _round_trips(redis) == 1
    assert prisma.db.mock_calls == []
    assert (user.user_id, team.team_id, membership.team_id, org.organization_id) == (USER_ID, TEAM_ID, TEAM_ID, ORG_ID)


@pytest.mark.asyncio
async def test_hot_regime_costs_nothing():
    redis = FakeAuthCache()
    prisma = _prisma()
    cache = _cache(redis)
    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma)
    redis.commands.clear()
    prisma.db.reset_mock()

    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma)
    await _read_all_through_getters(cache, prisma)

    assert _round_trips(redis) == 0
    assert prisma.db.mock_calls == []


@pytest.mark.asyncio
async def test_partial_redis_hit_queries_only_the_missing_objects():
    seeded = FakeAuthCache()
    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=_prisma())
    for key in (f"team_id:{TEAM_ID}", f"{TEAM_ID}_{USER_ID}", f"team_membership:{USER_ID}:{TEAM_ID}"):
        del seeded.redis[key]

    prisma = _prisma()
    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=prisma)

    assert prisma.db.query_first.await_count == 1
    assert prisma.db.query_first.await_args.args[1:] == (None, TEAM_ID, USER_ID, None, None)

    del seeded.redis[f"org_id:{ORG_ID}:with_budget"]
    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=prisma)
    assert prisma.db.query_first.await_args.args[1:] == (None, None, None, ORG_ID, None)


@pytest.mark.asyncio
async def test_deleted_team_cache_entry_is_refetched_and_the_update_is_visible():
    redis = FakeAuthCache()
    prisma = _prisma()
    cache = _cache(redis)
    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma)

    await cache.async_delete_cache(f"team_id:{TEAM_ID}")
    prisma.db.query_first.return_value = {**ALL_ROWS, "team_row": {**TEAM_ROW, "blocked": True}}
    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma)

    team = await get_team_object(team_id=TEAM_ID, prisma_client=prisma, user_api_key_cache=cache)
    assert team.blocked is True
    assert prisma.db.query_first.await_count == 2
    assert prisma.db.query_first.await_args.args[1:] == (None, TEAM_ID, None, None, None)


@pytest.mark.asyncio
async def test_row_missing_a_required_column_is_not_cached_so_the_getter_still_fails_closed():
    prisma = _prisma({**ALL_ROWS, "team_row": {"max_budget": 1.0}})
    prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
    cache = _cache(FakeAuthCache())

    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma)

    with pytest.raises(HTTPException) as exc:
        await get_team_object(team_id=TEAM_ID, prisma_client=prisma, user_api_key_cache=cache)
    assert exc.value.status_code == 404
    assert prisma.db.litellm_teamtable.find_unique.await_count == 1
    assert cache.get_local(USER_ID) is not None


@pytest.mark.asyncio
async def test_absent_rows_are_not_cached_as_present():
    redis = FakeAuthCache()
    prisma = _prisma({key: None for key in ALL_ROWS})
    cache = _cache(redis)

    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma)

    assert _writes(redis) == []
    assert cache.get_local(f"team_id:{TEAM_ID}") is None


@pytest.mark.asyncio
async def test_a_failed_redis_read_reads_as_a_miss_and_the_database_fill_takes_over():
    prisma = _prisma()
    cache = _cache(FakeAuthCache(fail=ConnectionError("redis down")))

    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma)

    assert prisma.db.query_first.await_count == 1
    assert cache.get_local(USER_ID)["user_id"] == USER_ID


@pytest.mark.asyncio
async def test_no_prisma_still_uses_redis_but_never_queries():
    seeded = FakeAuthCache()
    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=_cache(seeded), prisma_client=_prisma())
    redis = FakeAuthCache(redis=seeded.redis)
    cache = _cache(redis)

    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=None)

    assert _round_trips(redis) == 1
    assert cache.get_local(f"org_id:{ORG_ID}") is not None


@pytest.mark.asyncio
async def test_no_redis_goes_straight_to_one_query():
    prisma = _prisma()
    cache = _cache(None)

    await prefetch_auth_objects(refs=_refs(), user_api_key_cache=cache, prisma_client=prisma)

    assert prisma.db.query_first.await_count == 1
    assert cache.get_local(f"team_membership:{USER_ID}:{TEAM_ID}") is not None


@pytest.mark.asyncio
async def test_identity_prefetch_warms_the_end_user_so_its_getter_needs_neither_redis_nor_the_database():
    end_user_key = end_user_cache_key("eu-1")
    redis = FakeAuthCache(redis={end_user_key: {"user_id": "eu-1", "blocked": False, "spend": 0.0}})
    cache = _cache(redis)
    prisma = _prisma()

    await prefetch_identity_keys([end_user_key, end_user_restricted_registry_cache_key()], cache)
    end_user = await get_end_user_object(end_user_id="eu-1", prisma_client=prisma, user_api_key_cache=cache)

    assert end_user is not None and end_user.user_id == "eu-1"
    assert redis.commands == [("MGET", (end_user_key, end_user_restricted_registry_cache_key()))]
    assert prisma.db.mock_calls == []


@pytest.mark.asyncio
async def test_identity_prefetch_does_not_cache_an_absent_entry_as_present():
    redis = FakeAuthCache()
    cache = _cache(redis)

    await prefetch_identity_keys([end_user_cache_key("eu-absent")], cache)

    assert _round_trips(redis) == 1
    assert cache.get_local(end_user_cache_key("eu-absent")) is None