diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 5d9e0377c98..c3d61943288 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -3193,6 +3193,88 @@ async def test_generate_key_foreign_org_without_team_still_enforces_membership() assert mock_validate_org.call_args.kwargs["organization_id"] == foreign_org_id +@pytest.mark.asyncio +async def test_generate_key_foreign_org_with_mismatched_team_still_enforces_membership(): + """VERIA-55: when a team is present but its organization_id differs from the + organization_id on the key request, the org-membership check must still run.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _common_key_generation_helper, + ) + + team_org_id = "other-org" + foreign_org_id = "someone-elses-org" + + mock_team_table = MagicMock() + mock_team_table.organization_id = team_org_id + mock_team_table.metadata = None + + mock_validate_org = AsyncMock() + mock_generate_key = AsyncMock( + return_value={ + "key": "sk-test-key", + "expires": None, + "user_id": "alice", + "team_id": "team-1", + } + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_mcp_servers_against_team", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_search_tools_against_team", + new_callable=AsyncMock, + ), + patch( + "litellm_enterprise.proxy.management_endpoints.key_management_endpoints.apply_enterprise_key_management_params", + side_effect=lambda data, team_table: data, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._validate_caller_can_assign_key_org", + mock_validate_org, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_org_object", + new_callable=AsyncMock, + return_value=MagicMock(litellm_budget_table=None), + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._check_org_key_limits", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + mock_generate_key, + ), + ): + await _common_key_generation_helper( + data=GenerateKeyRequest( + user_id="alice", + team_id="team-1", + organization_id=foreign_org_id, + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ), + litellm_changed_by=None, + team_table=mock_team_table, + ) + + mock_validate_org.assert_awaited_once() + assert mock_validate_org.call_args.kwargs["organization_id"] == foreign_org_id + + # ============================================ # Organization Key Limit Tests # ============================================