diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 8e2358c99..a0e020b11 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -312,7 +312,7 @@ async def user_info( try: if prisma_client is None: raise Exception( - f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" + "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) ## GET USER ROW ## if user_id is not None: @@ -365,7 +365,14 @@ async def user_info( getattr(caller_user_info, "user_role", None) == LitellmUserRoles.PROXY_ADMIN ): - teams_2 = await prisma_client.db.litellm_teamtable.find_many() + from litellm.proxy.management_endpoints.team_endpoints import list_team + + teams_2 = await list_team( + http_request=Request( + scope={"type": "http", "path": "/user/info"}, + ), + user_api_key_dict=user_api_key_dict, + ) else: teams_2 = await prisma_client.get_data( team_id_list=caller_user_info.teams, diff --git a/litellm/tests/test_proxy_server.py b/litellm/tests/test_proxy_server.py index dee20a273..757eef6d6 100644 --- a/litellm/tests/test_proxy_server.py +++ b/litellm/tests/test_proxy_server.py @@ -928,3 +928,41 @@ async def test_create_team_member_add(prisma_client, new_member_method): mock_client.call_args.kwargs["data"]["create"]["budget_duration"] == litellm.internal_user_budget_duration ) + + +@pytest.mark.asyncio +async def test_user_info_team_list(prisma_client): + """Assert user_info for admin calls team_list function""" + from litellm.proxy._types import LiteLLM_UserTable + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + await litellm.proxy.proxy_server.prisma_client.connect() + + from litellm.proxy.management_endpoints.internal_user_endpoints import user_info + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.list_team", + new_callable=AsyncMock, + ) as mock_client: + + prisma_client.get_data = AsyncMock( + return_value=LiteLLM_UserTable( + user_role="proxy_admin", + user_id="default_user_id", + max_budget=None, + user_email="", + ) + ) + + try: + await user_info( + user_id=None, + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-1234", user_id="default_user_id" + ), + ) + except Exception: + pass + + mock_client.assert_called() diff --git a/tests/test_users.py b/tests/test_users.py index 632dd8f36..8113fd080 100644 --- a/tests/test_users.py +++ b/tests/test_users.py @@ -99,13 +99,6 @@ async def test_user_info(): ) assert status == 403 - ## check if returned teams as admin == all teams ## - admin_info = await get_user_info( - session=session, get_user="", call_user="sk-1234", view_all=True - ) - all_teams = await list_teams(session=session, i=0) - assert len(admin_info["teams"]) == len(all_teams) - @pytest.mark.asyncio async def test_user_update():