mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-12 12:22:36 +00:00
fix(proxy): enforce organization boundaries in admin operations
Validate org admin role against all requested organizations instead of returning on first match. Scope team list queries to the caller's permitted organizations when filtering by user_id.
This commit is contained in:
user
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
Reference in New Issue
Block a user