mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix: document org scope behavior, fix test mocks, add org admin tests
- Document intentional legacy-matching behavior: when user_id is provided to an org admin, no org filter is applied (returns all of that user's teams across all orgs, same as legacy endpoint) - Fix two existing security tests to properly patch user_api_key_cache, proxy_logging_obj, and get_user_object instead of relying on incidental error handling - Add three new org admin test cases: - Org admin sees org-scoped teams (200 with correct where clause) - Org admin rejected when filtering by other org (403) - Org admin with user_id filter returns target user's teams Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
yuneng-jiang
parent
0485a1859a
commit
41a7747e8c
2 changed files with 280 additions and 8 deletions
|
|
@ -3272,7 +3272,12 @@ 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
|
||||
# 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.
|
||||
where_conditions["organization_id"] = {"in": org_admin_org_ids}
|
||||
|
||||
if user_id:
|
||||
|
|
|
|||
|
|
@ -2062,7 +2062,14 @@ 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:
|
||||
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
|
||||
|
||||
# Should raise HTTPException with 401 status
|
||||
|
|
@ -2103,7 +2110,14 @@ 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:
|
||||
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
|
||||
|
||||
# Should raise HTTPException with 401 status
|
||||
|
|
@ -2300,27 +2314,280 @@ async def test_list_team_v2_with_status_deleted():
|
|||
assert len(result["teams"]) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_team_v2_org_admin_sees_org_teams():
|
||||
"""
|
||||
Test that an org admin (internal_user role with org_admin membership)
|
||||
can list teams scoped to their organisations without getting a 401.
|
||||
"""
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_OrganizationMembershipTable,
|
||||
LiteLLM_UserTable,
|
||||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.team_endpoints import list_team_v2
|
||||
|
||||
mock_request = Mock(spec=Request)
|
||||
mock_user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="org_admin_user",
|
||||
)
|
||||
|
||||
mock_user = LiteLLM_UserTable(
|
||||
user_id="org_admin_user",
|
||||
teams=[],
|
||||
organization_memberships=[
|
||||
LiteLLM_OrganizationMembershipTable(
|
||||
user_id="org_admin_user",
|
||||
organization_id="org_A",
|
||||
user_role="org_admin",
|
||||
spend=0.0,
|
||||
created_at=datetime.now(),
|
||||
updated_at=datetime.now(),
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
mock_team = Mock()
|
||||
mock_team.model_dump.return_value = {
|
||||
"team_id": "team_in_org_A",
|
||||
"team_alias": "Org A Team",
|
||||
"organization_id": "org_A",
|
||||
"members_with_roles": [{"user_id": "u1", "role": "user"}],
|
||||
}
|
||||
mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team])
|
||||
mock_db.litellm_teamtable.count = AsyncMock(return_value=1)
|
||||
|
||||
result = await list_team_v2(
|
||||
http_request=mock_request,
|
||||
user_id=None,
|
||||
organization_id=None,
|
||||
team_id=None,
|
||||
team_alias=None,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
page=1,
|
||||
page_size=10,
|
||||
sort_by=None,
|
||||
sort_order="asc",
|
||||
status=None,
|
||||
)
|
||||
|
||||
assert result["total"] == 1
|
||||
assert len(result["teams"]) == 1
|
||||
assert result["teams"][0].members_count == 1
|
||||
|
||||
# Verify org-scoped where clause
|
||||
where = mock_db.litellm_teamtable.find_many.call_args.kwargs["where"]
|
||||
assert where["organization_id"] == {"in": ["org_A"]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_team_v2_org_admin_cannot_view_other_orgs():
|
||||
"""
|
||||
Test that an org admin is rejected with 403 when filtering by an
|
||||
organisation they do not administer.
|
||||
"""
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
from fastapi import HTTPException, Request
|
||||
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_OrganizationMembershipTable,
|
||||
LiteLLM_UserTable,
|
||||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.team_endpoints import list_team_v2
|
||||
|
||||
mock_request = Mock(spec=Request)
|
||||
mock_user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="org_admin_user",
|
||||
)
|
||||
|
||||
mock_user = LiteLLM_UserTable(
|
||||
user_id="org_admin_user",
|
||||
teams=[],
|
||||
organization_memberships=[
|
||||
LiteLLM_OrganizationMembershipTable(
|
||||
user_id="org_admin_user",
|
||||
organization_id="org_A",
|
||||
user_role="org_admin",
|
||||
spend=0.0,
|
||||
created_at=datetime.now(),
|
||||
updated_at=datetime.now(),
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
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()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await list_team_v2(
|
||||
http_request=mock_request,
|
||||
user_id=None,
|
||||
organization_id="org_B", # not their org
|
||||
team_id=None,
|
||||
team_alias=None,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
page=1,
|
||||
page_size=10,
|
||||
sort_by=None,
|
||||
sort_order="asc",
|
||||
status=None,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "only view teams within your organizations" in str(
|
||||
exc_info.value.detail
|
||||
).lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_team_v2_org_admin_with_user_id_returns_user_teams():
|
||||
"""
|
||||
Test that an org admin passing user_id gets that user's direct team
|
||||
memberships (not all org teams).
|
||||
"""
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_OrganizationMembershipTable,
|
||||
LiteLLM_UserTable,
|
||||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.team_endpoints import list_team_v2
|
||||
|
||||
mock_request = Mock(spec=Request)
|
||||
mock_user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
user_id="org_admin_user",
|
||||
)
|
||||
|
||||
mock_org_admin = LiteLLM_UserTable(
|
||||
user_id="org_admin_user",
|
||||
teams=["team_1"],
|
||||
organization_memberships=[
|
||||
LiteLLM_OrganizationMembershipTable(
|
||||
user_id="org_admin_user",
|
||||
organization_id="org_A",
|
||||
user_role="org_admin",
|
||||
spend=0.0,
|
||||
created_at=datetime.now(),
|
||||
updated_at=datetime.now(),
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
# The target user whose teams we want to list
|
||||
mock_target_user = LiteLLM_UserTable(
|
||||
user_id="target_user",
|
||||
teams=["team_X", "team_Y"],
|
||||
)
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def mock_get_user_object(**kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
# First call: org admin lookup in list_team_v2
|
||||
# Second call: target user lookup in _build_team_list_where_conditions
|
||||
if call_count == 1:
|
||||
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,
|
||||
):
|
||||
mock_db = Mock()
|
||||
mock_prisma.db = mock_db
|
||||
|
||||
mock_team = Mock()
|
||||
mock_team.model_dump.return_value = {
|
||||
"team_id": "team_X",
|
||||
"team_alias": "Target Team",
|
||||
"members_with_roles": [{"user_id": "target_user", "role": "user"}],
|
||||
}
|
||||
mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team])
|
||||
mock_db.litellm_teamtable.count = AsyncMock(return_value=1)
|
||||
|
||||
result = await list_team_v2(
|
||||
http_request=mock_request,
|
||||
user_id="target_user",
|
||||
organization_id=None,
|
||||
team_id=None,
|
||||
team_alias=None,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
page=1,
|
||||
page_size=10,
|
||||
sort_by=None,
|
||||
sort_order="asc",
|
||||
status=None,
|
||||
)
|
||||
|
||||
assert result["total"] == 1
|
||||
|
||||
# Verify the where clause filters by user's teams, not 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
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_team_v2_with_invalid_status():
|
||||
"""
|
||||
Test that invalid status parameter raises HTTPException.
|
||||
"""
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
|
||||
from fastapi import HTTPException, Request
|
||||
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.team_endpoints import list_team_v2
|
||||
|
||||
|
||||
# Mock request
|
||||
mock_request = Mock(spec=Request)
|
||||
|
||||
|
||||
# Mock admin user
|
||||
mock_user_api_key_dict_admin = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_id="admin_user_123",
|
||||
)
|
||||
|
||||
|
||||
mock_prisma_client = Mock()
|
||||
|
||||
# Mock prisma_client to be non-None
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue