diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 04554aeb322..32795a1874d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -13,7 +13,7 @@ import re 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 @@ -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( @@ -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 @@ -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." + }, ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index cff4cf48fc4..f9257ffe078 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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()) @@ -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()) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 9ae916db0a7..6dad7cb08d0 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -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", @@ -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() diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 7d4a406c99a..057b56ce317 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -705,3 +705,30 @@ async def test_get_tag_objects_batch(): 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)