diff --git a/tests/test_litellm/proxy/management_endpoints/management_v1/test_users.py b/tests/test_litellm/proxy/management_endpoints/management_v1/test_users.py index 2e66729194c..06cca8e376f 100644 --- a/tests/test_litellm/proxy/management_endpoints/management_v1/test_users.py +++ b/tests/test_litellm/proxy/management_endpoints/management_v1/test_users.py @@ -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 diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index dbebf460d3e..d8aea3268c8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -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(