From 41a7747e8cc99b133dd72ac6de6f55aaa9708e0d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 17 Mar 2026 21:56:57 -0700 Subject: [PATCH] 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) --- .../management_endpoints/team_endpoints.py | 7 +- .../test_team_endpoints.py | 281 +++++++++++++++++- 2 files changed, 280 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 531cd675892..e7cd6c1748f 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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: 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 122ac749f2b..c325b0b6fce 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -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