From de84b2edce05c225e7903d914040b5da4e508a25 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 16 Jan 2026 19:24:51 -0800 Subject: [PATCH] Fixing tests and adding proper returns --- litellm/proxy/_types.py | 5 +- .../key_management_endpoints.py | 20 ++- .../management_endpoints/team_endpoints.py | 133 ++++++++++----- .../management_endpoints/team_endpoints.py | 5 +- .../test_key_management_endpoints.py | 152 ++++++++++-------- 5 files changed, 199 insertions(+), 116 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0f2cd4b3a7..f717cca9b9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2229,6 +2229,9 @@ class UserAPIKeyAuth( @model_validator(mode="before") @classmethod def check_api_key(cls, values): + # If values is already an instance (not a dict), return it as-is + if not isinstance(values, dict): + return values if values.get("api_key") is not None: values.update( {"token": cls._safe_hash_litellm_api_key(values.get("api_key"))} @@ -3359,7 +3362,7 @@ class TeamListResponseObject(LiteLLM_TeamTable): class KeyListResponseObject(TypedDict, total=False): - keys: List[Union[str, UserAPIKeyAuth]] + keys: List[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]] total_count: Optional[int] current_page: Optional[int] total_pages: Optional[int] diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index bebe30b0a4..e40a44edf5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3543,19 +3543,31 @@ async def _list_key_helper( user_map = {user.user_id: user for user in users} # Prepare response - key_list: List[Union[str, UserAPIKeyAuth]] = [] + key_list: List[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]] = [] for key in keys: - key_dict = key.dict() + # Convert Prisma model to dict (supports both Pydantic v1 and v2) + try: + key_dict = key.model_dump() + except Exception: + # Fallback for Pydantic v1 compatibility + key_dict = key.dict() # Attach object_permission if object_permission_id is set (only for non-deleted keys) if not use_deleted_table: key_dict = await attach_object_permission_to_dict(key_dict, prisma_client) # Include user information if expand includes "user" if expand and "user" in expand and key.user_id and key.user_id in user_map: - key_dict["user"] = user_map[key.user_id].dict() + try: + key_dict["user"] = user_map[key.user_id].model_dump() + except Exception: + key_dict["user"] = user_map[key.user_id].dict() if return_full_object is True or (expand and "user" in expand): - key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object + if use_deleted_table: + # Use deleted key type to preserve deleted_at, deleted_by, etc. + key_list.append(LiteLLM_DeletedVerificationToken(**key_dict)) + else: + key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object else: _token = key_dict.get("token") key_list.append(cast(str, _token)) # Return only the token diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 24d7639af3..b082c3ea55 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2955,6 +2955,83 @@ async def list_available_teams( return available_teams_correct_type +async def _build_team_list_where_conditions( + prisma_client: PrismaClient, + team_id: Optional[str], + team_alias: Optional[str], + organization_id: Optional[str], + user_id: Optional[str], + use_deleted_table: bool, +) -> Dict[str, Any]: + """Build where conditions for team list query.""" + where_conditions: Dict[str, Any] = {} + + if team_id: + where_conditions["team_id"] = team_id + + if team_alias: + where_conditions["team_alias"] = { + "contains": team_alias, + "mode": "insensitive", # Case-insensitive search + } + + if organization_id: + where_conditions["organization_id"] = organization_id + + if user_id: + try: + user_object = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_id} + ) + except Exception: + raise HTTPException( + status_code=404, + detail={"error": f"User not found, passed user_id={user_id}"}, + ) + if user_object is None: + raise HTTPException( + status_code=404, + detail={"error": f"User not found, passed user_id={user_id}"}, + ) + user_object_correct_type = LiteLLM_UserTable(**user_object.model_dump()) + + if use_deleted_table: + where_conditions["members"] = {"has": user_id} + else: + if team_id is None: + where_conditions["team_id"] = {"in": user_object_correct_type.teams} + elif team_id in user_object_correct_type.teams: + where_conditions["team_id"] = team_id + else: + raise HTTPException( + status_code=404, + detail={"error": f"User is not a member of team_id={team_id}"}, + ) + + return where_conditions + + +def _convert_teams_to_response( + teams: List[Any], use_deleted_table: bool +) -> List[Union[LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]]: + """Convert Prisma models to Pydantic models.""" + team_list = [] + if teams: + for team in teams: + # Convert Prisma model to dict (supports both Pydantic v1 and v2) + try: + team_dict = team.model_dump() + except Exception: + # Fallback for Pydantic v1 compatibility + team_dict = team.dict() + if use_deleted_table: + # Use deleted team type to preserve deleted_at, deleted_by, etc. + team_list.append(LiteLLM_DeletedTeamTable(**team_dict)) + else: + team_list.append(LiteLLM_TeamTable(**team_dict)) + return team_list + + @router.get( "/v2/team/list", tags=["team management"], @@ -3056,49 +3133,14 @@ async def list_team_v2( skip = (page - 1) * page_size # Build where conditions based on provided parameters - where_conditions: Dict[str, Any] = {} - - if team_id: - where_conditions["team_id"] = team_id - - if team_alias: - where_conditions["team_alias"] = { - "contains": team_alias, - "mode": "insensitive", # Case-insensitive search - } - - if organization_id: - where_conditions["organization_id"] = organization_id - - if user_id: - try: - user_object = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_id} - ) - except Exception: - raise HTTPException( - status_code=404, - detail={"error": f"User not found, passed user_id={user_id}"}, - ) - if user_object is None: - raise HTTPException( - status_code=404, - detail={"error": f"User not found, passed user_id={user_id}"}, - ) - user_object_correct_type = LiteLLM_UserTable(**user_object.model_dump()) - - if use_deleted_table: - where_conditions["members"] = {"has": user_id} - else: - if team_id is None: - where_conditions["team_id"] = {"in": user_object_correct_type.teams} - elif team_id in user_object_correct_type.teams: - where_conditions["team_id"] = team_id - else: - raise HTTPException( - status_code=404, - detail={"error": f"User is not a member of team_id={team_id}"}, - ) + where_conditions = await _build_team_list_where_conditions( + prisma_client=prisma_client, + team_id=team_id, + team_alias=team_alias, + organization_id=organization_id, + user_id=user_id, + use_deleted_table=use_deleted_table, + ) # Build order_by conditions valid_sort_columns = ["team_id", "team_alias", "created_at"] @@ -3133,8 +3175,11 @@ async def list_team_v2( # Calculate total pages total_pages = -(-total_count // page_size) # Ceiling division + # Convert Prisma models to Pydantic models, preserving deleted fields when applicable + team_list = _convert_teams_to_response(teams, use_deleted_table) + return { - "teams": [team.model_dump() for team in teams] if teams else [], + "teams": team_list, "total": total_count, "page": page, "page_size": page_size, diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 957e5d60eb..77816fa78c 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -1,8 +1,9 @@ -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Union from pydantic import BaseModel from litellm.proxy._types import ( + LiteLLM_DeletedTeamTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_UserTable, @@ -45,7 +46,7 @@ class UpdateTeamMemberPermissionsRequest(BaseModel): class TeamListResponse(BaseModel): """Response to get the list of teams""" - teams: List[LiteLLM_TeamTable] + teams: List[Union[LiteLLM_TeamTable, LiteLLM_DeletedTeamTable]] total: int page: int page_size: int diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 613ff31c1e..7d31f76209 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -10,7 +10,7 @@ sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch from fastapi import HTTPException @@ -3800,47 +3800,61 @@ async def test_list_keys_with_expand_user(): mock_prisma_client = AsyncMock() # Create mock keys with user_ids - mock_key1 = MagicMock() - mock_key1.token = "token1" - mock_key1.user_id = "user123" - mock_key1.dict.return_value = { + key1_dict = { "token": "token1", "user_id": "user123", "key_alias": "key1", "models": ["gpt-4"], } + mock_key1 = MagicMock() + mock_key1.token = "token1" + mock_key1.user_id = "user123" + # Set up model_dump() to raise AttributeError so it falls back to dict() + mock_key1.model_dump = MagicMock(side_effect=AttributeError("model_dump not available")) + mock_key1.dict = MagicMock(return_value=key1_dict) - mock_key2 = MagicMock() - mock_key2.token = "token2" - mock_key2.user_id = "user456" - mock_key2.dict.return_value = { + key2_dict = { "token": "token2", "user_id": "user456", "key_alias": "key2", "models": ["gpt-3.5-turbo"], } + mock_key2 = MagicMock() + mock_key2.token = "token2" + mock_key2.user_id = "user456" + # Set up model_dump() to raise AttributeError so it falls back to dict() + mock_key2.model_dump = MagicMock(side_effect=AttributeError("model_dump not available")) + mock_key2.dict = MagicMock(return_value=key2_dict) mock_find_many_keys = AsyncMock(return_value=[mock_key1, mock_key2]) mock_count_keys = AsyncMock(return_value=2) # Create mock users - mock_user1 = MagicMock() - mock_user1.user_id = "user123" - mock_user1.user_email = "user1@example.com" - mock_user1.dict.return_value = { + user1_dict = { "user_id": "user123", "user_email": "user1@example.com", "user_alias": "User One", } + mock_user1 = MagicMock() + # Set user_id as a real attribute (not a MagicMock) + mock_user1.user_id = "user123" + mock_user1.user_email = "user1@example.com" + # Set up both model_dump() and dict() to return the same dict + mock_user1.model_dump = MagicMock(return_value=user1_dict) + mock_user1.dict = MagicMock(return_value=user1_dict) - mock_user2 = MagicMock() - mock_user2.user_id = "user456" - mock_user2.user_email = "user2@example.com" - mock_user2.dict.return_value = { + user2_dict = { "user_id": "user456", "user_email": "user2@example.com", "user_alias": "User Two", } + mock_user2 = MagicMock() + # Set user_id as a real attribute (not a MagicMock) + mock_user2.user_id = "user456" + mock_user2.user_email = "user2@example.com" + # Set up both model_dump() and dict() to return the same dict + mock_user2.model_dump = MagicMock(return_value=user2_dict) + mock_user2.dict = MagicMock(return_value=user2_dict) mock_find_many_users = AsyncMock(return_value=[mock_user1, mock_user2]) @@ -3848,60 +3862,68 @@ async def test_list_keys_with_expand_user(): mock_prisma_client.db.litellm_verificationtoken.count = mock_count_keys mock_prisma_client.db.litellm_usertable.find_many = mock_find_many_users - args = { - "prisma_client": mock_prisma_client, - "page": 1, - "size": 50, - "user_id": None, - "team_id": None, - "organization_id": None, - "key_alias": None, - "key_hash": None, - "exclude_team_id": None, - "return_full_object": False, # This should be overridden by expand=user - "admin_team_ids": None, - "include_created_by_keys": False, - "expand": ["user"], # Test the expand parameter - } + # Patch attach_object_permission_to_dict to just return the dict unchanged + async def mock_attach_object_permission(d, _): + return d + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.attach_object_permission_to_dict", + side_effect=mock_attach_object_permission, + ): + args = { + "prisma_client": mock_prisma_client, + "page": 1, + "size": 50, + "user_id": None, + "team_id": None, + "organization_id": None, + "key_alias": None, + "key_hash": None, + "exclude_team_id": None, + "return_full_object": False, # This should be overridden by expand=user + "admin_team_ids": None, + "include_created_by_keys": False, + "expand": ["user"], # Test the expand parameter + } - result = await _list_key_helper(**args) + result = await _list_key_helper(**args) - # Verify that keys were fetched - mock_find_many_keys.assert_called_once() - mock_count_keys.assert_called_once() + # Verify that keys were fetched + mock_find_many_keys.assert_called_once() + mock_count_keys.assert_called_once() - # Verify that users were fetched - # Note: Order doesn't matter for the 'in' query, so we just check that both user_ids are present - call_args = mock_find_many_users.call_args - assert call_args is not None - where_clause = call_args.kwargs["where"] - assert "user_id" in where_clause - assert "in" in where_clause["user_id"] - user_ids_in_query = set(where_clause["user_id"]["in"]) - assert user_ids_in_query == {"user123", "user456"} + # Verify that users were fetched + # Note: Order doesn't matter for the 'in' query, so we just check that both user_ids are present + call_args = mock_find_many_users.call_args + assert call_args is not None + where_clause = call_args.kwargs["where"] + assert "user_id" in where_clause + assert "in" in where_clause["user_id"] + user_ids_in_query = set(where_clause["user_id"]["in"]) + assert user_ids_in_query == {"user123", "user456"} - # Verify response structure - assert len(result["keys"]) == 2 - assert result["total_count"] == 2 - assert result["current_page"] == 1 - assert result["total_pages"] == 1 + # Verify response structure + assert len(result["keys"]) == 2 + assert result["total_count"] == 2 + assert result["current_page"] == 1 + assert result["total_pages"] == 1 - # Verify that user data is included in the response - # Since expand=user is specified, keys should be full objects - assert isinstance(result["keys"][0], UserAPIKeyAuth) - assert isinstance(result["keys"][1], UserAPIKeyAuth) + # Verify that user data is included in the response + # Since expand=user is specified, keys should be full objects + assert isinstance(result["keys"][0], UserAPIKeyAuth) + assert isinstance(result["keys"][1], UserAPIKeyAuth) - # Verify user data is attached to keys - assert result["keys"][0].user == { - "user_id": "user123", - "user_email": "user1@example.com", - "user_alias": "User One", - } - assert result["keys"][1].user == { - "user_id": "user456", - "user_email": "user2@example.com", - "user_alias": "User Two", - } + # Verify user data is attached to keys + assert result["keys"][0].user == { + "user_id": "user123", + "user_email": "user1@example.com", + "user_alias": "User One", + } + assert result["keys"][1].user == { + "user_id": "user456", + "user_email": "user2@example.com", + "user_alias": "User Two", + } @pytest.mark.asyncio