diff --git a/litellm/proxy/auth/auth_checks_organization.py b/litellm/proxy/auth/auth_checks_organization.py index 50efe13720..d89afcffa9 100644 --- a/litellm/proxy/auth/auth_checks_organization.py +++ b/litellm/proxy/auth/auth_checks_organization.py @@ -144,7 +144,7 @@ def _user_is_org_admin( user_object: Optional[LiteLLM_UserTable] = None, ) -> bool: """ - Helper function to check if user is an org admin for any of the passed organizations. + Helper function to check if user is an org admin for all of the passed organizations. Checks both: - `organization_id` (singular string) — legacy callers @@ -168,9 +168,13 @@ def _user_is_org_admin( if not candidate_org_ids: return False - for _membership in user_object.organization_memberships: - if _membership.organization_id in candidate_org_ids: - if _membership.user_role == LitellmUserRoles.ORG_ADMIN.value: - return True + # Build set of orgs where user is admin + admin_org_ids = { + _membership.organization_id + for _membership in user_object.organization_memberships + if _membership.user_role == LitellmUserRoles.ORG_ADMIN.value + and _membership.organization_id is not None + } - return False + # User must be admin of ALL requested orgs, not just any one + return all(org_id in admin_org_ids for org_id in candidate_org_ids) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 138469312e..edd92cb83d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -120,6 +120,7 @@ def _sanitize_for_log(value: Any) -> str: text = repr(value) return text.replace("\r", "").replace("\n", "") + async def _verify_team_access( team_obj: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth, @@ -314,10 +315,8 @@ class TeamMemberBudgetHandler: return # Batch-fetch existing memberships for this team (avoids N+1 queries) - existing_memberships = ( - await prisma_client.db.litellm_teammembership.find_many( - where={"team_id": team_id} - ) + existing_memberships = await prisma_client.db.litellm_teammembership.find_many( + where={"team_id": team_id} ) existing_user_ids = {m.user_id for m in existing_memberships} @@ -1659,12 +1658,12 @@ async def update_team( # noqa: PLR0915 updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Optional[ - LiteLLM_TeamTable - ] = await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, - data=updated_kv, - include={"litellm_model_table": True}, # type: ignore + team_row: Optional[LiteLLM_TeamTable] = ( + await prisma_client.db.litellm_teamtable.update( + where={"team_id": data.team_id}, + data=updated_kv, + include={"litellm_model_table": True}, # type: ignore + ) ) if team_row is None or team_row.team_id is None: @@ -2411,13 +2410,13 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={ - "user_id": {"in": list(user_ids_to_delete)}, - "team_id": data.team_id, - } + keys_to_delete: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={ + "user_id": {"in": list(user_ids_to_delete)}, + "team_id": data.team_id, + } + ) ) if keys_to_delete: @@ -2801,10 +2800,10 @@ async def delete_team( team_rows: List[LiteLLM_TeamTable] = [] for team_id in data.team_ids: try: - team_row_base: Optional[ - BaseModel - ] = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} + team_row_base: Optional[BaseModel] = ( + await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) ) if team_row_base is None: raise Exception @@ -2870,10 +2869,10 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": {"in": data.team_ids}} + keys_to_delete: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={"team_id": {"in": data.team_ids}} + ) ) if keys_to_delete: @@ -3110,11 +3109,11 @@ async def team_info( ) try: - team_info: Optional[ - BaseModel - ] = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id}, - include={"object_permission": True}, + team_info: Optional[BaseModel] = ( + await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + include={"object_permission": True}, + ) ) if team_info is None: raise Exception @@ -3405,6 +3404,7 @@ async def _get_org_admin_org_ids( m.organization_id for m in (caller_user.organization_memberships or []) if m.user_role == LitellmUserRoles.ORG_ADMIN.value + and m.organization_id is not None ] return org_ids if org_ids else None @@ -3439,13 +3439,8 @@ async def _build_team_list_where_conditions( if organization_id: where_conditions["organization_id"] = organization_id - elif org_admin_org_ids is not None and not user_id: - # Org admin without explicit org or user filter: scope to their orgs. - # NOTE: when user_id is provided, no org filter is applied — the - # query returns all teams the target user belongs to across all - # organisations. This matches the legacy /team/list behaviour in - # _authorize_and_filter_teams which fetches direct-membership teams - # without an org constraint. + elif org_admin_org_ids is not None: + # Org admin: always scope to their orgs, even when filtering by user_id. where_conditions["organization_id"] = {"in": org_admin_org_ids} if user_id: @@ -3815,7 +3810,7 @@ async def _authorize_and_filter_teams( Authorize the /team/list request and return filtered teams. - Proxy admins: all teams (or filtered by user_id if provided). - - Org admins: teams from their orgs + teams they are direct members of. + - Org admins: teams from their orgs (scoped to user_id if provided). - Own query (user_id matches caller): teams the user is a member of. - Others: 401. """ @@ -3843,6 +3838,7 @@ async def _authorize_and_filter_teams( m.organization_id for m in (caller_user.organization_memberships or []) if m.user_role == LitellmUserRoles.ORG_ADMIN.value + and m.organization_id is not None ] if not allowed_org_ids: allowed_org_ids = None @@ -3865,20 +3861,13 @@ async def _authorize_and_filter_teams( ) if not user_id: return list(org_teams) - # Also include teams the user is a direct member of (outside their orgs) - seen_team_ids = {team.team_id for team in org_teams} - all_teams = list(org_teams) - # Prisma doesn't support filtering JSON array fields, so we fetch by membership separately - member_teams = await prisma_client.db.litellm_teamtable.find_many( - where={"team_id": {"not_in": list(seen_team_ids)}} if seen_team_ids else {}, - include={"litellm_model_table": True}, - ) - for team in member_teams: - if team.members_with_roles and any( - m.get("user_id") == user_id for m in team.members_with_roles - ): - all_teams.append(team) - return all_teams + # Filter org teams to only those where the target user is a member + return [ + team + for team in org_teams + if team.members_with_roles + and any(m.get("user_id") == user_id for m in team.members_with_roles) + ] elif user_id: # Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays) response = await prisma_client.db.litellm_teamtable.find_many( diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index bee6642dec..9b4bd79049 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1010,12 +1010,15 @@ async def test_validate_team_member_add_permissions_non_admin(): team.organization_id = None # Mock the helper functions to return False - with patch( - "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", - return_value=False, - ), patch( - "litellm.proxy.management_endpoints.team_endpoints._is_available_team", - return_value=False, + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=False, + ), ): # Should raise HTTPException for non-admin with pytest.raises(HTTPException) as exc_info: @@ -1257,19 +1260,17 @@ async def test_update_team_team_member_budget_not_passed_to_db(): user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( - "litellm.proxy.proxy_server.llm_router" - ) as mock_llm_router, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.proxy_logging_obj" - ) as mock_logging, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.auth.auth_checks._cache_team_object" - ) as mock_cache_team, patch( - "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" - ) as mock_upsert_budget: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch("litellm.proxy.proxy_server.llm_router") as mock_llm_router, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch("litellm.proxy.auth.auth_checks._cache_team_object") as mock_cache_team, + patch( + "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" + ) as mock_upsert_budget, + ): # Setup mock prisma client mock_existing_team = MagicMock() mock_existing_team.model_dump.return_value = { @@ -1690,19 +1691,17 @@ async def test_update_team_with_team_member_budget_duration(): user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( - "litellm.proxy.proxy_server.llm_router" - ) as mock_llm_router, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.proxy_logging_obj" - ) as mock_logging, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.auth.auth_checks._cache_team_object" - ) as mock_cache_team, patch( - "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" - ) as mock_upsert_budget: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch("litellm.proxy.proxy_server.llm_router") as mock_llm_router, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch("litellm.proxy.auth.auth_checks._cache_team_object") as mock_cache_team, + patch( + "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" + ) as mock_upsert_budget, + ): mock_existing_team = MagicMock() mock_existing_team.model_dump.return_value = { "team_id": "test_team_id", @@ -1777,7 +1776,9 @@ async def test_backfill_team_member_budget_entries_creates_missing_memberships() from unittest.mock import AsyncMock, MagicMock from litellm.proxy._types import Member - from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) team_id = "team-abc" budget_id = "budget-xyz" @@ -1847,7 +1848,9 @@ async def test_backfill_team_member_budget_entries_no_op_when_all_exist(): from unittest.mock import AsyncMock, MagicMock from litellm.proxy._types import Member - from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) team_id = "team-abc" budget_id = "budget-xyz" @@ -1886,7 +1889,9 @@ async def test_backfill_team_member_budget_entries_empty_members(): """ from unittest.mock import AsyncMock, MagicMock - from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) mock_prisma = MagicMock() mock_prisma.db.litellm_teammembership.find_many = AsyncMock(return_value=[]) @@ -2092,11 +2097,14 @@ async def test_bulk_team_member_add_all_users_flag(): updated_team_memberships=[], ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.management_endpoints.team_endpoints.team_member_add", - new_callable=AsyncMock, - return_value=mock_team_response, - ) as mock_team_member_add: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch( + "litellm.proxy.management_endpoints.team_endpoints.team_member_add", + new_callable=AsyncMock, + return_value=mock_team_response, + ) as mock_team_member_add, + ): # Mock the database find_many call mock_prisma.db.litellm_usertable.find_many = AsyncMock( return_value=mock_db_users @@ -2213,12 +2221,15 @@ async def test_list_team_v2_security_check_non_admin_user(): user_id="non_admin_user_123", ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ), patch("litellm.proxy.proxy_server.proxy_logging_obj"), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_user_object", - new_callable=AsyncMock, - return_value=None, + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), ): mock_prisma_client.return_value = MagicMock() # Mock non-None prisma client @@ -2260,12 +2271,15 @@ async def test_list_team_v2_security_check_non_admin_user_other_user(): user_id="non_admin_user_123", ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ), patch("litellm.proxy.proxy_server.proxy_logging_obj"), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_user_object", - new_callable=AsyncMock, - return_value=None, + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=None, + ), ): mock_prisma_client.return_value = MagicMock() # Mock non-None prisma client @@ -2305,9 +2319,11 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): user_id="non_admin_user_123", ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ), patch("litellm.proxy.proxy_server.proxy_logging_obj"): + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + ): # Mock prisma client and database operations mock_db = Mock() mock_prisma_client.db = mock_db @@ -2509,12 +2525,15 @@ async def test_list_team_v2_org_admin_sees_org_teams(): ], ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ), patch("litellm.proxy.proxy_server.proxy_logging_obj"), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_user_object", - new_callable=AsyncMock, - return_value=mock_user, + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=mock_user, + ), ): mock_db = Mock() mock_prisma.db = mock_db @@ -2592,12 +2611,15 @@ async def test_list_team_v2_org_admin_cannot_view_other_orgs(): ], ) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ), patch("litellm.proxy.proxy_server.proxy_logging_obj"), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_user_object", - new_callable=AsyncMock, - return_value=mock_user, + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + return_value=mock_user, + ), ): mock_prisma.db = Mock() @@ -2680,11 +2702,14 @@ async def test_list_team_v2_org_admin_with_user_id_returns_user_teams(): return mock_org_admin return mock_target_user - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ), patch("litellm.proxy.proxy_server.proxy_logging_obj"), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_user_object", - side_effect=mock_get_user_object, + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + side_effect=mock_get_user_object, + ), ): mock_db = Mock() mock_prisma.db = mock_db @@ -2714,10 +2739,10 @@ async def test_list_team_v2_org_admin_with_user_id_returns_user_teams(): assert result["total"] == 1 - # Verify the where clause filters by user's teams, not org scope + # Verify the where clause filters by user's teams AND org scope where = mock_db.litellm_teamtable.find_many.call_args.kwargs["where"] assert where["team_id"] == {"in": ["team_X", "team_Y"]} - assert "organization_id" not in where + assert where["organization_id"] == {"in": ["org_A"]} @pytest.mark.asyncio @@ -2913,15 +2938,15 @@ async def test_new_team_max_budget_exceeds_user_max_budget(): dummy_request = MagicMock(spec=Request) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server._license_check" - ) as mock_license, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + ): # Setup basic mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_license.is_team_count_over_limit.return_value = False @@ -2982,15 +3007,15 @@ async def test_new_team_max_budget_within_user_limit(): dummy_request = MagicMock(spec=Request) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server._license_check" - ) as mock_license, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + ): # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_license.is_team_count_over_limit.return_value = False @@ -3111,17 +3136,18 @@ async def test_new_team_org_scoped_budget_bypasses_user_limit(): dummy_request = MagicMock(spec=Request) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server._license_check" - ) as mock_license, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object" - ) as mock_get_org: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object" + ) as mock_get_org, + ): # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_license.is_team_count_over_limit.return_value = False @@ -3253,17 +3279,18 @@ async def test_new_team_org_scoped_models_bypasses_user_limit(): dummy_request = MagicMock(spec=Request) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server._license_check" - ) as mock_license, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object" - ) as mock_get_org: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object" + ) as mock_get_org, + ): # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_license.is_team_count_over_limit.return_value = False @@ -3393,13 +3420,14 @@ async def test_new_team_standalone_validates_against_user_models(monkeypatch): dummy_request = MagicMock(spec=Request) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server._license_check" - ) as mock_license, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + ): # Setup basic mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_license.is_team_count_over_limit.return_value = False @@ -3460,15 +3488,15 @@ async def test_new_team_standalone_validates_against_user_budget(): dummy_request = MagicMock(spec=Request) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server._license_check" - ) as mock_license, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + ): # Setup basic mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_license.is_team_count_over_limit.return_value = False @@ -3534,17 +3562,18 @@ async def test_new_team_org_scoped_budget_exceeds_org_limit(): dummy_request = MagicMock(spec=Request) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server._license_check" - ) as mock_license, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object" - ) as mock_get_org: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object" + ) as mock_get_org, + ): # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_license.is_team_count_over_limit.return_value = False @@ -3613,17 +3642,18 @@ async def test_new_team_org_scoped_models_not_in_org_models(): dummy_request = MagicMock(spec=Request) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server._license_check" - ) as mock_license, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object" - ) as mock_get_org: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object" + ) as mock_get_org, + ): # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_license.is_team_count_over_limit.return_value = False @@ -3688,13 +3718,14 @@ async def test_update_team_standalone_budget_exceeds_user_limit(): dummy_request = MagicMock(spec=Request) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + ): # Mock existing standalone team (no organization_id) mock_existing_team = MagicMock() mock_existing_team.team_id = "standalone-team-123" @@ -3778,16 +3809,18 @@ async def test_update_team_org_scoped_budget_exceeds_org_limit(): mock_org.models = ["gpt-4"] mock_org.litellm_budget_table = mock_budget_table - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object", - new=AsyncMock(return_value=mock_org), - ) as mock_get_org: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org), + ) as mock_get_org, + ): # Mock existing org-scoped team mock_existing_team = MagicMock() mock_existing_team.team_id = "org-team-456" @@ -3852,13 +3885,14 @@ async def test_update_team_standalone_models_exceeds_user_limit(): dummy_request = MagicMock(spec=Request) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + ): # Mock existing standalone team (no organization_id) mock_existing_team = MagicMock() mock_existing_team.team_id = "standalone-team-models-123" @@ -3936,16 +3970,18 @@ async def test_update_team_org_scoped_budget_bypasses_user_limit(): mock_org.models = ["gpt-4", "gpt-3.5-turbo"] mock_org.litellm_budget_table = mock_budget_table - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object", - new=AsyncMock(return_value=mock_org), - ) as mock_get_org: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org), + ) as mock_get_org, + ): # Mock existing org-scoped team mock_existing_team = MagicMock() mock_existing_team.team_id = "org-team-update-budget-123" @@ -4044,16 +4080,18 @@ async def test_update_team_org_scoped_models_bypasses_user_limit(): mock_org.models = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus"] mock_org.litellm_budget_table = None - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object", - new=AsyncMock(return_value=mock_org), - ) as mock_get_org: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org), + ) as mock_get_org, + ): # Mock existing org-scoped team mock_existing_team = MagicMock() mock_existing_team.team_id = "org-team-update-models-123" @@ -4145,16 +4183,18 @@ async def test_update_team_org_scoped_models_not_in_org_models(): mock_org.models = ["gpt-4", "gpt-3.5-turbo"] # claude-3-opus is NOT allowed mock_org.litellm_budget_table = None - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object", - new=AsyncMock(return_value=mock_org), - ) as mock_get_org: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org), + ) as mock_get_org, + ): # Mock existing org-scoped team mock_existing_team = MagicMock() mock_existing_team.team_id = "org-team-update-models-fail-123" @@ -4231,16 +4271,18 @@ async def test_update_team_org_scoped_models_with_all_proxy_models(): mock_org.models = [SpecialModelNames.all_proxy_models.value] # Allows all models mock_org.litellm_budget_table = None - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit, patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object", - new=AsyncMock(return_value=mock_org), - ) as mock_get_org: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org), + ) as mock_get_org, + ): # Mock existing org-scoped team mock_existing_team = MagicMock() mock_existing_team.team_id = "org-team-all-proxy-models-123" @@ -4333,10 +4375,10 @@ async def test_update_team_tpm_limit_exceeds_user_limit(): dummy_request = MagicMock(spec=Request) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), ): # Mock existing standalone team mock_existing_team = MagicMock() @@ -4397,10 +4439,10 @@ async def test_update_team_rpm_limit_exceeds_user_limit(): dummy_request = MagicMock(spec=Request) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), ): # Mock existing standalone team mock_existing_team = MagicMock() @@ -4479,15 +4521,15 @@ async def test_new_team_org_scoped_tpm_exceeds_org_limit(): mock_org.models = ["gpt-4"] mock_org.litellm_budget_table = mock_budget_table - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server._license_check" - ) as mock_license, patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object", - new=AsyncMock(return_value=mock_org), + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org), + ), ): mock_license.is_team_count_over_limit.return_value = False mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) @@ -4555,15 +4597,15 @@ async def test_new_team_org_scoped_rpm_exceeds_org_limit(): mock_org.models = ["gpt-4"] mock_org.litellm_budget_table = mock_budget_table - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server._license_check" - ) as mock_license, patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object", - new=AsyncMock(return_value=mock_org), + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org), + ), ): mock_license.is_team_count_over_limit.return_value = False mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) @@ -4634,20 +4676,22 @@ async def test_new_team_org_scoped_tpm_rpm_bypasses_user_limit(): mock_org.models = ["gpt-4"] mock_org.litellm_budget_table = mock_budget_table - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server._license_check" - ) as mock_license, patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object", - new=AsyncMock(return_value=mock_org), - ), patch( - "litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", - new=AsyncMock(), + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org), + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", + new=AsyncMock(), + ), ): mock_license.is_team_count_over_limit.return_value = False mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) @@ -4736,13 +4780,14 @@ async def test_update_team_org_scoped_tpm_exceeds_org_limit(): mock_org.models = ["gpt-4"] mock_org.litellm_budget_table = mock_budget_table - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object", - new=AsyncMock(return_value=mock_org), + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org), + ), ): # Mock existing org-scoped team mock_existing_team = MagicMock() @@ -4822,13 +4867,14 @@ async def test_update_team_org_scoped_rpm_exceeds_org_limit(): mock_org.models = ["gpt-4"] mock_org.litellm_budget_table = mock_budget_table - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object", - new=AsyncMock(return_value=mock_org), + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org), + ), ): # Mock existing org-scoped team mock_existing_team = MagicMock() @@ -4911,15 +4957,15 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit(): mock_org.models = ["gpt-4"] mock_org.litellm_budget_table = mock_budget_table - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.proxy_logging_obj" - ) as mock_logging, patch( - "litellm.proxy.management_endpoints.team_endpoints.get_org_object", - new=AsyncMock(return_value=mock_org), + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_org_object", + new=AsyncMock(return_value=mock_org), + ), ): # Mock existing org-scoped team mock_existing_team = MagicMock() @@ -5036,17 +5082,18 @@ async def test_update_team_guardrails_with_org_id(): "teams": [], } - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ), patch( - "litellm.proxy.proxy_server.premium_user", - True, # Required for guardrails feature - ), patch( - "litellm.proxy.proxy_server.llm_router", MagicMock() + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ), + patch( + "litellm.proxy.proxy_server.premium_user", + True, # Required for guardrails feature + ), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), ): # Mock existing team - must have compatible models with organization mock_existing_team = MagicMock() @@ -5601,15 +5648,15 @@ async def test_new_team_soft_budget_validation( dummy_request = MagicMock(spec=Request) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server._license_check" - ) as mock_license, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server._license_check") as mock_license, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + ): # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_license.is_team_count_over_limit.return_value = False @@ -5799,13 +5846,14 @@ async def test_update_team_soft_budget_validation( dummy_request = MagicMock(spec=Request) - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( - "litellm.proxy.proxy_server.user_api_key_cache" - ) as mock_cache, patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() - ) as mock_audit: + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + ): # Mock existing team with existing budgets mock_existing_team = MagicMock() mock_existing_team.team_id = "test-team-123" @@ -6794,14 +6842,18 @@ async def test_list_team_v1_batches_key_queries(): key3 = MagicMock() key3.team_id = "team-2" - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( - "litellm.proxy.management_endpoints.team_endpoints._authorize_and_filter_teams", - new_callable=AsyncMock, - return_value=[team1, team2], - ), patch( - "litellm.proxy.management_endpoints.team_endpoints.get_all_team_memberships", - new_callable=AsyncMock, - return_value=[], + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch( + "litellm.proxy.management_endpoints.team_endpoints._authorize_and_filter_teams", + new_callable=AsyncMock, + return_value=[team1, team2], + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_all_team_memberships", + new_callable=AsyncMock, + return_value=[], + ), ): async def filtered_find_many(**kwargs): @@ -7070,16 +7122,17 @@ async def test_update_team_rejects_unauthorized_caller(): from litellm.proxy._types import UpdateTeamRequest - with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( - "litellm.proxy.proxy_server.llm_router" - ), patch("litellm.proxy.proxy_server.user_api_key_cache"), patch( - "litellm.proxy.proxy_server.proxy_logging_obj" - ), patch( - "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" - ), patch( - "litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", - new_callable=AsyncMock, - return_value=False, + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, + patch("litellm.proxy.proxy_server.llm_router"), + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_org_admin_for_team", + new_callable=AsyncMock, + return_value=False, + ), ): mock_existing_team = MagicMock() mock_existing_team.model_dump.return_value = {