mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
780c504e00
* test: migrate tests/proxy_behavior into integration and unit tiers (partial) * test(integration): run settings-dependent management tests on an owned proxy * test(integration): assert stored values in the rewritten management rows * test(integration): pin sibling team allocations against the org guaranteed throughput cap * test(integration): keep Prisma clients out of the shared integration conftest import path --------- Co-authored-by: yuneng <yuneng@berri.ai>
190 lines
10 KiB
Python
190 lines
10 KiB
Python
"""Runs the auth prefetch's raw SQL against a real Postgres: the join must bind the membership to the requested
|
|
team and hand the getters rows they validate. The per-regime round-trip counts are unit-tested with fakes in
|
|
tests/unit/proxy/auth/test_auth_object_prefetch.py."""
|
|
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
|
|
from litellm.caching.in_memory_cache import InMemoryCache
|
|
from litellm.proxy.auth.auth_checks import (
|
|
get_org_object,
|
|
get_team_membership,
|
|
get_team_object,
|
|
get_user_object,
|
|
)
|
|
from litellm.proxy.auth.auth_object_prefetch import AuthObjectRefs, prefetch_auth_objects
|
|
from litellm.proxy.auth.team_grants import team_model_aliases
|
|
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
|
from litellm.proxy.utils import PrismaClient
|
|
from tests.integration._support.prisma_clients import connected_proxy_prisma
|
|
|
|
pytestmark = pytest.mark.asyncio(loop_scope="session")
|
|
|
|
|
|
@pytest_asyncio.fixture(scope="module", loop_scope="session")
|
|
async def prisma() -> AsyncIterator[PrismaClient]:
|
|
async with connected_proxy_prisma() as client:
|
|
yield client
|
|
|
|
|
|
def _frozen_cache() -> UserApiKeyCache:
|
|
"""The org entries carry a 5s TTL; a frozen clock keeps a slow first call from expiring them mid-test."""
|
|
return UserApiKeyCache(in_memory_cache=InMemoryCache(clock=lambda: 1_000_000.0), redis_cache=None)
|
|
|
|
|
|
def _dead_db() -> MagicMock:
|
|
prisma = MagicMock(name="prisma_client")
|
|
prisma.db.query_first = AsyncMock(return_value=None)
|
|
return prisma
|
|
|
|
|
|
async def test_join_binds_the_membership_to_the_requested_team(prisma):
|
|
"""A user in two teams with different member budgets must get the requested team's row."""
|
|
run = uuid4().hex
|
|
user_id, team_a, team_b, org_id = (f"pf-user-{run}", f"pf-team-a-{run}", f"pf-team-b-{run}", f"pf-org-{run}")
|
|
try:
|
|
await prisma.db.litellm_budgettable.create(
|
|
data={"budget_id": f"a-{run}", "max_budget": 11.0, "created_by": "t", "updated_by": "t"}
|
|
)
|
|
await prisma.db.litellm_budgettable.create(
|
|
data={"budget_id": f"b-{run}", "max_budget": 22.0, "created_by": "t", "updated_by": "t"}
|
|
)
|
|
await prisma.db.litellm_organizationtable.create(
|
|
data={
|
|
"organization_id": org_id,
|
|
"organization_alias": "pf",
|
|
"created_by": "t",
|
|
"updated_by": "t",
|
|
"litellm_budget_table": {"connect": {"budget_id": f"b-{run}"}},
|
|
}
|
|
)
|
|
await prisma.db.litellm_usertable.create(data={"user_id": user_id, "max_budget": 33.0})
|
|
await prisma.db.litellm_teamtable.create(data={"team_id": team_a, "organization_id": org_id, "max_budget": 1.0})
|
|
await prisma.db.litellm_teamtable.create(data={"team_id": team_b, "max_budget": 2.0})
|
|
await prisma.db.litellm_teammembership.create(
|
|
data={"user_id": user_id, "team_id": team_a, "litellm_budget_table": {"connect": {"budget_id": f"a-{run}"}}}
|
|
)
|
|
await prisma.db.litellm_teammembership.create(
|
|
data={"user_id": user_id, "team_id": team_b, "litellm_budget_table": {"connect": {"budget_id": f"b-{run}"}}}
|
|
)
|
|
|
|
cache = _frozen_cache()
|
|
refs = AuthObjectRefs(user_id=user_id, team_id=team_a, membership_user_id=user_id, organization_id=org_id)
|
|
await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma)
|
|
assert cache.in_memory_cache.get_cache(f"org_id:{org_id}") is not None
|
|
|
|
dead_db = _dead_db()
|
|
membership = await get_team_membership(
|
|
user_id=user_id, team_id=team_a, prisma_client=dead_db, user_api_key_cache=cache
|
|
)
|
|
team = await get_team_object(team_id=team_a, prisma_client=dead_db, user_api_key_cache=cache)
|
|
user = await get_user_object(
|
|
user_id=user_id, prisma_client=dead_db, user_api_key_cache=cache, user_id_upsert=False
|
|
)
|
|
org = await get_org_object(org_id=org_id, prisma_client=dead_db, user_api_key_cache=cache)
|
|
assert dead_db.db.mock_calls == [], "getters must be served from the prefetched cache"
|
|
|
|
assert membership is not None and membership.litellm_budget_table is not None
|
|
assert (membership.team_id, membership.litellm_budget_table.max_budget) == (team_a, 11.0)
|
|
assert (team.team_id, team.max_budget, team.organization_id, team.models) == (team_a, 1.0, org_id, [])
|
|
assert user is not None and user.max_budget == 33.0
|
|
assert org is not None and (org.organization_id, org.models) == (org_id, [])
|
|
finally:
|
|
await prisma.db.litellm_teammembership.delete_many(where={"user_id": user_id})
|
|
await prisma.db.litellm_teamtable.delete_many(where={"team_id": {"in": [team_a, team_b]}})
|
|
await prisma.db.litellm_usertable.delete_many(where={"user_id": user_id})
|
|
await prisma.db.litellm_organizationtable.delete_many(where={"organization_id": org_id})
|
|
await prisma.db.litellm_budgettable.delete_many(where={"budget_id": {"in": [f"a-{run}", f"b-{run}"]}})
|
|
|
|
|
|
async def test_join_reads_team_model_aliases_from_the_mapped_column(prisma):
|
|
"""The model table stores aliases in a column named ``aliases``; the cached team must expose ``model_aliases``."""
|
|
run = uuid4().hex
|
|
team_id = f"pf-team-{run}"
|
|
aliases = {"gpt-4o": f"gpt-4o-{run}"}
|
|
model_table = await prisma.db.litellm_modeltable.create(
|
|
data={"model_aliases": json.dumps(aliases), "created_by": "t", "updated_by": "t"}
|
|
)
|
|
try:
|
|
await prisma.db.litellm_teamtable.create(data={"team_id": team_id, "model_id": model_table.id})
|
|
expected_team = await prisma.db.litellm_teamtable.find_unique(
|
|
where={"team_id": team_id}, include={"litellm_model_table": True}
|
|
)
|
|
|
|
cache = _frozen_cache()
|
|
refs = AuthObjectRefs(user_id=None, team_id=team_id, membership_user_id=None, organization_id=None)
|
|
await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma)
|
|
|
|
dead_db = _dead_db()
|
|
team = await get_team_object(team_id=team_id, prisma_client=dead_db, user_api_key_cache=cache)
|
|
assert dead_db.db.mock_calls == [], "getters must be served from the prefetched cache"
|
|
|
|
assert expected_team is not None and expected_team.litellm_model_table is not None
|
|
assert team.litellm_model_table is not None
|
|
assert team.litellm_model_table.model_aliases == expected_team.litellm_model_table.model_aliases == aliases
|
|
assert team_model_aliases(team) == aliases
|
|
finally:
|
|
await prisma.db.litellm_teamtable.delete_many(where={"team_id": team_id})
|
|
await prisma.db.litellm_modeltable.delete_many(where={"id": model_table.id})
|
|
|
|
|
|
async def test_join_reads_null_nested_lists_the_way_prisma_does(prisma):
|
|
"""Prisma reads a NULL scalar list as []; the nested permission and budget rows must match, not carry null."""
|
|
run = uuid4().hex
|
|
user_id, team_id, permission_id, budget_id = (f"pf-user-{run}", f"pf-team-{run}", f"pf-perm-{run}", f"pf-bud-{run}")
|
|
try:
|
|
await prisma.db.litellm_objectpermissiontable.create(data={"object_permission_id": permission_id})
|
|
await prisma.db.litellm_budgettable.create(data={"budget_id": budget_id, "created_by": "t", "updated_by": "t"})
|
|
await prisma.db.execute_raw(
|
|
'UPDATE "LiteLLM_ObjectPermissionTable" SET mcp_servers = NULL, models = NULL '
|
|
"WHERE object_permission_id = $1",
|
|
permission_id,
|
|
)
|
|
await prisma.db.execute_raw(
|
|
'UPDATE "LiteLLM_BudgetTable" SET allowed_models = NULL WHERE budget_id = $1', budget_id
|
|
)
|
|
await prisma.db.litellm_usertable.create(data={"user_id": user_id})
|
|
await prisma.db.litellm_teamtable.create(data={"team_id": team_id, "object_permission_id": permission_id})
|
|
await prisma.db.litellm_teammembership.create(
|
|
data={"user_id": user_id, "team_id": team_id, "litellm_budget_table": {"connect": {"budget_id": budget_id}}}
|
|
)
|
|
expected_team = await prisma.db.litellm_teamtable.find_unique(
|
|
where={"team_id": team_id}, include={"object_permission": True}
|
|
)
|
|
expected_membership = await prisma.db.litellm_teammembership.find_unique(
|
|
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True}
|
|
)
|
|
|
|
cache = _frozen_cache()
|
|
refs = AuthObjectRefs(user_id=user_id, team_id=team_id, membership_user_id=user_id, organization_id=None)
|
|
await prefetch_auth_objects(refs=refs, user_api_key_cache=cache, prisma_client=prisma)
|
|
|
|
dead_db = _dead_db()
|
|
team = await get_team_object(team_id=team_id, prisma_client=dead_db, user_api_key_cache=cache)
|
|
membership = await get_team_membership(
|
|
user_id=user_id, team_id=team_id, prisma_client=dead_db, user_api_key_cache=cache
|
|
)
|
|
assert dead_db.db.mock_calls == [], "getters must be served from the prefetched cache"
|
|
|
|
assert expected_team is not None and expected_team.object_permission is not None
|
|
assert team.object_permission is not None
|
|
assert team.object_permission.mcp_servers == expected_team.object_permission.mcp_servers == []
|
|
assert team.object_permission.models == expected_team.object_permission.models == []
|
|
assert expected_membership is not None and expected_membership.litellm_budget_table is not None
|
|
assert membership is not None and membership.litellm_budget_table is not None
|
|
assert (
|
|
membership.litellm_budget_table.allowed_models
|
|
== expected_membership.litellm_budget_table.allowed_models
|
|
== []
|
|
)
|
|
finally:
|
|
await prisma.db.litellm_teammembership.delete_many(where={"user_id": user_id})
|
|
await prisma.db.litellm_teamtable.delete_many(where={"team_id": team_id})
|
|
await prisma.db.litellm_usertable.delete_many(where={"user_id": user_id})
|
|
await prisma.db.litellm_objectpermissiontable.delete_many(where={"object_permission_id": permission_id})
|
|
await prisma.db.litellm_budgettable.delete_many(where={"budget_id": budget_id})
|