Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 12 additions & 6 deletions litellm/proxy/auth/auth_checks.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
import time
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast

from fastapi import Request, status
from fastapi import HTTPException, Request, status
from pydantic import BaseModel

import litellm
Expand Down Expand Up @@ -1274,7 +1274,7 @@ async def get_team_object(
- if not, then raise an error

Raises:
- Exception: If team doesn't exist in db or cache
- HTTPException: If team doesn't exist in db or cache (status_code=404)
"""
if prisma_client is None:
raise Exception(
Expand All @@ -1296,8 +1296,11 @@ async def get_team_object(
return cached_team_obj

if check_cache_only:
raise Exception(
f"Team doesn't exist in cache + check_cache_only=True. Team={team_id}."
raise HTTPException(
status_code=404,
detail={
"error": f"Team doesn't exist in cache + check_cache_only=True. Team={team_id}."
},
)

# else, check db
Expand All @@ -1313,8 +1316,11 @@ async def get_team_object(
team_id_upsert=team_id_upsert,
)
except Exception:
raise Exception(
f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call."
raise HTTPException(
status_code=404,
detail={
"error": f"Team doesn't exist in db. Team={team_id}. Create team via `/team/new` call."
},
)


Expand Down
10 changes: 0 additions & 10 deletions litellm/proxy/management_endpoints/team_endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -3247,11 +3247,6 @@ async def team_member_permissions(
check_cache_only=False,
check_db_only=True,
)
if existing_team_row is None:
raise HTTPException(
status_code=404,
detail={"error": f"Team not found for team_id={team_id}"},
)

complete_team_data = LiteLLM_TeamTable(**existing_team_row.model_dump())

Expand Down Expand Up @@ -3320,11 +3315,6 @@ async def update_team_member_permissions(
check_cache_only=False,
check_db_only=True,
)
if existing_team_row is None:
raise HTTPException(
status_code=404,
detail={"error": f"Team not found for team_id={data.team_id}"},
)

complete_team_data = LiteLLM_TeamTable(**existing_team_row.model_dump())

Expand Down
4 changes: 3 additions & 1 deletion tests/proxy_unit_tests/test_proxy_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -953,6 +953,8 @@ async def test_get_team_redis(client_no_auth):

redis_cache = RedisCache()

from fastapi import HTTPException

with patch.object(
redis_cache,
"async_get_cache",
Expand All @@ -966,7 +968,7 @@ async def test_get_team_redis(client_no_auth):
proxy_logging_obj=proxy_logging_obj,
prisma_client=AsyncMock(),
)
except Exception as e:
except HTTPException:
pass

mock_client.assert_called_once()
Expand Down
27 changes: 27 additions & 0 deletions tests/test_litellm/proxy/auth/test_auth_checks.py
Original file line number Diff line number Diff line change
Expand Up @@ -705,3 +705,30 @@ async def mock_get_cache(key):
assert "tag:uncached-1" in cached_keys
assert "tag:uncached-2" in cached_keys
assert "tag:uncached-3" in cached_keys


@pytest.mark.asyncio
async def test_get_team_object_raises_404_when_not_found():
from litellm.proxy.auth.auth_checks import get_team_object
from fastapi import HTTPException
from unittest.mock import AsyncMock, MagicMock

mock_prisma_client = MagicMock()
mock_db = AsyncMock()
mock_prisma_client.db = mock_db
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)

mock_cache = MagicMock()
mock_cache.async_get_cache = AsyncMock(return_value=None)

with pytest.raises(HTTPException) as exc_info:
await get_team_object(
team_id="nonexistent-team",
prisma_client=mock_prisma_client,
user_api_key_cache=mock_cache,
check_cache_only=False,
check_db_only=True,
)

assert exc_info.value.status_code == 404
assert "Team doesn't exist in db" in str(exc_info.value.detail)
Loading