mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
test: patch the prisma client once per user test instead of once per collaborator
The management_v1 fixture reached into UserRepository.table and the spend-counter helper to stand the route up. Seeding prisma_client.db is enough, since the repository reads the table off it, and the other two patches were doing nothing: the admin-name patch set the value the constant already had, and the counter only fires when the payload carries spend. The user_info_v2 tests each built their own prisma double, so a fixture now seeds the rows the group reads and the tests just say which users exist.
This commit is contained in:
ryan-crabbe-berri
parent
128b74b5a3
commit
423563ed9c
2 changed files with 74 additions and 71 deletions
|
|
@ -107,21 +107,20 @@ def caller():
|
|||
|
||||
@pytest.fixture
|
||||
def prisma(mocker):
|
||||
"""A prisma double whose `update_data` call is what the assertions inspect."""
|
||||
"""A prisma double whose `update_data` call is what the assertions inspect.
|
||||
|
||||
`UserRepository.table` reads straight off `prisma_client.db`, so seeding that is enough and the
|
||||
repository itself stays real. The module global is the one seam the route has: it imports
|
||||
`prisma_client` from `proxy_server` at call time.
|
||||
"""
|
||||
table = MagicMock()
|
||||
table.find_first = AsyncMock(return_value=_user_row())
|
||||
client = MagicMock()
|
||||
client.update_data = AsyncMock(return_value={"user_id": TARGET_ID, "data": _user_row()})
|
||||
client.get_data = AsyncMock(return_value=[])
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", client)
|
||||
mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id")
|
||||
mocker.patch("litellm.proxy.proxy_server._invalidate_spend_counter", AsyncMock())
|
||||
table = MagicMock()
|
||||
table.find_first = AsyncMock(return_value=_user_row())
|
||||
mocker.patch(
|
||||
"litellm.repositories.user_repository.UserRepository.table",
|
||||
new_callable=mocker.PropertyMock,
|
||||
return_value=table,
|
||||
)
|
||||
client.db.litellm_usertable = table
|
||||
client.table = table
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", client)
|
||||
return client
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2589,8 +2589,32 @@ async def test_user_update_rejects_silent_create_for_non_proxy_admin(mocker):
|
|||
# =====================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def user_rows(mocker):
|
||||
"""Seed the user table `/v2/user/info` reads, keyed by user_id, and return nothing for the rest.
|
||||
|
||||
One seam for the whole `user_info_v2` group, so a test says which rows exist and nothing else.
|
||||
"""
|
||||
|
||||
def _seed(**rows):
|
||||
prisma_client = mocker.MagicMock()
|
||||
|
||||
async def find_unique(*args, **kwargs):
|
||||
row = rows.get(kwargs.get("where", {}).get("user_id"))
|
||||
if row is None:
|
||||
return None
|
||||
stub = mocker.MagicMock()
|
||||
stub.model_dump.return_value = row
|
||||
return stub
|
||||
|
||||
prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(side_effect=find_unique)
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client)
|
||||
|
||||
return _seed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_info_v2_proxy_admin_can_query_any_user(mocker):
|
||||
async def test_user_info_v2_proxy_admin_can_query_any_user(mocker, user_rows):
|
||||
"""
|
||||
Test that proxy admin can query any user via /v2/user/info.
|
||||
"""
|
||||
|
|
@ -2599,37 +2623,27 @@ async def test_user_info_v2_proxy_admin_can_query_any_user(mocker):
|
|||
from litellm.proxy._types import UserInfoV2Response
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2
|
||||
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
|
||||
mock_user_row = mocker.MagicMock()
|
||||
mock_user_row.model_dump.return_value = {
|
||||
"user_id": "target-user-123",
|
||||
"user_email": "target@example.com",
|
||||
"user_alias": "Target User",
|
||||
"user_role": "internal_user",
|
||||
"spend": 42.5,
|
||||
"max_budget": 100.0,
|
||||
"models": ["gpt-4"],
|
||||
"budget_duration": "30d",
|
||||
"budget_reset_at": None,
|
||||
"metadata": {"team": "engineering"},
|
||||
"created_at": datetime(2024, 1, 1, tzinfo=timezone.utc),
|
||||
"updated_at": datetime(2024, 6, 1, tzinfo=timezone.utc),
|
||||
"sso_user_id": "sso-abc",
|
||||
"teams": ["team-1", "team-2"],
|
||||
}
|
||||
|
||||
async def mock_find_unique(*args, **kwargs):
|
||||
if kwargs.get("where", {}).get("user_id") == "target-user-123":
|
||||
return mock_user_row
|
||||
return None
|
||||
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(
|
||||
side_effect=mock_find_unique
|
||||
user_rows(
|
||||
**{
|
||||
"target-user-123": {
|
||||
"user_id": "target-user-123",
|
||||
"user_email": "target@example.com",
|
||||
"user_alias": "Target User",
|
||||
"user_role": "internal_user",
|
||||
"spend": 42.5,
|
||||
"max_budget": 100.0,
|
||||
"models": ["gpt-4"],
|
||||
"budget_duration": "30d",
|
||||
"budget_reset_at": None,
|
||||
"metadata": {"team": "engineering"},
|
||||
"created_at": datetime(2024, 1, 1, tzinfo=timezone.utc),
|
||||
"updated_at": datetime(2024, 6, 1, tzinfo=timezone.utc),
|
||||
"sso_user_id": "sso-abc",
|
||||
"teams": ["team-1", "team-2"],
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
||||
mock_request = mocker.MagicMock(spec=Request)
|
||||
|
||||
admin_key = UserAPIKeyAuth(
|
||||
|
|
@ -2744,7 +2758,7 @@ def test_build_user_info_response_redacts_scim_enterprise_metadata():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_info_v2_internal_user_can_query_self(mocker):
|
||||
async def test_user_info_v2_internal_user_can_query_self(mocker, user_rows):
|
||||
"""
|
||||
Test that an internal user can query their own info.
|
||||
"""
|
||||
|
|
@ -2753,37 +2767,27 @@ async def test_user_info_v2_internal_user_can_query_self(mocker):
|
|||
from litellm.proxy._types import UserInfoV2Response
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2
|
||||
|
||||
mock_prisma_client = mocker.MagicMock()
|
||||
|
||||
mock_user_row = mocker.MagicMock()
|
||||
mock_user_row.model_dump.return_value = {
|
||||
"user_id": "self-user",
|
||||
"user_email": "self@example.com",
|
||||
"user_alias": None,
|
||||
"user_role": "internal_user",
|
||||
"spend": 10.0,
|
||||
"max_budget": None,
|
||||
"models": [],
|
||||
"budget_duration": None,
|
||||
"budget_reset_at": None,
|
||||
"metadata": None,
|
||||
"created_at": None,
|
||||
"updated_at": None,
|
||||
"sso_user_id": None,
|
||||
"teams": [],
|
||||
}
|
||||
|
||||
async def mock_find_unique(*args, **kwargs):
|
||||
if kwargs.get("where", {}).get("user_id") == "self-user":
|
||||
return mock_user_row
|
||||
return None
|
||||
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(
|
||||
side_effect=mock_find_unique
|
||||
user_rows(
|
||||
**{
|
||||
"self-user": {
|
||||
"user_id": "self-user",
|
||||
"user_email": "self@example.com",
|
||||
"user_alias": None,
|
||||
"user_role": "internal_user",
|
||||
"spend": 10.0,
|
||||
"max_budget": None,
|
||||
"models": [],
|
||||
"budget_duration": None,
|
||||
"budget_reset_at": None,
|
||||
"metadata": None,
|
||||
"created_at": None,
|
||||
"updated_at": None,
|
||||
"sso_user_id": None,
|
||||
"teams": [],
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
||||
mock_request = mocker.MagicMock(spec=Request)
|
||||
|
||||
user_key = UserAPIKeyAuth(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue