diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py index f4e91157f21..67e9dda4cef 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_manager_role.py @@ -262,10 +262,6 @@ class TestCreateMcpServerAsManager: reload_servers_from_database=AsyncMock(), ), ), - patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_team_object", - AsyncMock(return_value=mock_team), - ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common", mock_handle_update, @@ -280,3 +276,80 @@ class TestCreateMcpServerAsManager: assert "existing_server" in mcp_servers assert "new_server_id" in mcp_servers assert call_kwargs["existing_object_permission_id"] == "perm1" + + async def test_create_links_new_permission_to_team_when_none_exists(self): + """When team has no object_permission_id, create should link the new one.""" + from litellm.proxy._types import NewMCPServerRequest + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + add_mcp_server, + ) + + user_auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="user1", + api_key="sk-test", + team_id="team1", + ) + + payload = NewMCPServerRequest( + server_name="test_server", + url="https://example.com/mcp", + team_id="team1", + ) + + # Team with NO object_permission_id + mock_team = LiteLLM_TeamTableCachedObj( + team_id="team1", + members_with_roles=[ + Member(user_id="user1", role="mcp_server_manager"), + ], + object_permission_id=None, + ) + mock_team.object_permission = None + + created_server = MagicMock() + created_server.server_id = "new_server_id" + created_server.credentials = None + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.update = AsyncMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_prisma, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._assert_can_manage_team_mcp_server", + AsyncMock(return_value=("team1", mock_team)), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", + AsyncMock(return_value=created_server), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + MagicMock( + add_server=AsyncMock(), + reload_servers_from_database=AsyncMock(), + ), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.handle_update_object_permission_common", + AsyncMock(return_value="new_perm_id"), + ), + ): + await add_mcp_server(payload=payload, user_api_key_dict=user_auth) + + # Verify the team was updated with the new object_permission_id + mock_prisma.db.litellm_teamtable.update.assert_called_once_with( + where={"team_id": "team1"}, + data={"object_permission_id": "new_perm_id"}, + )