tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 12:23:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixing tests and adding proper returns

Browse Source
This commit is contained in:
yuneng-jiang
2026-01-16 19:24:51 -08:00
parent 6e8dd06d18
commit de84b2edce
5 changed files with 199 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_types.py
+4 -1
View File
@@ -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]
litellm/proxy/management_endpoints/key_management_endpoints.py
+16 -4
View File
@@ -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
litellm/proxy/management_endpoints/team_endpoints.py
+89 -44
View File
@@ -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,
litellm/types/proxy/management_endpoints/team_endpoints.py
+3 -2
View File
@@ -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
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
+87 -65
View File
@@ -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