From bdab813f21c5b4fa551eb198f463f39cea045363 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Sat, 21 Mar 2026 12:28:24 -0700
Subject: [PATCH] test: add MCP server team ownership tests

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../test_mcp_team_ownership.py                | 82 +++++++++++++++++++
 1 file changed, 82 insertions(+)
 create mode 100644 tests/test_litellm/proxy/management_endpoints/test_mcp_team_ownership.py

diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_team_ownership.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_team_ownership.py
new file mode 100644
index 00000000000..452eb4bc34f
--- /dev/null
+++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_team_ownership.py
@@ -0,0 +1,82 @@
+"""Tests for MCP server team ownership via MCPServerTable.team_id."""
+
+import uuid
+
+import pytest
+
+
+def test_prepare_mcp_server_data_includes_team_id():
+    """team_id should be included in data dict when set."""
+    from litellm.proxy._experimental.mcp_server.db import _prepare_mcp_server_data
+    from litellm.proxy._types import NewMCPServerRequest
+
+    team_id = str(uuid.uuid4())
+    request = NewMCPServerRequest(
+        server_name="test_server",
+        transport="http",
+        url="https://example.com/mcp",
+        team_id=team_id,
+    )
+    data_dict = _prepare_mcp_server_data(request)
+    assert data_dict["team_id"] == team_id
+
+
+def test_prepare_mcp_server_data_excludes_none_team_id():
+    """team_id=None should not be in the data dict (exclude_none=True)."""
+    from litellm.proxy._experimental.mcp_server.db import _prepare_mcp_server_data
+    from litellm.proxy._types import NewMCPServerRequest
+
+    request = NewMCPServerRequest(
+        server_name="test_server",
+        transport="http",
+        url="https://example.com/mcp",
+    )
+    data_dict = _prepare_mcp_server_data(request)
+    assert "team_id" not in data_dict
+
+
+def test_prepare_mcp_server_data_update_includes_team_id():
+    """UpdateMCPServerRequest with team_id should include it in data dict."""
+    from litellm.proxy._experimental.mcp_server.db import _prepare_mcp_server_data
+    from litellm.proxy._types import UpdateMCPServerRequest
+
+    team_id = str(uuid.uuid4())
+    request = UpdateMCPServerRequest(
+        server_id="test-server-id",
+        transport="http",
+        url="https://example.com/mcp",
+        team_id=team_id,
+    )
+    data_dict = _prepare_mcp_server_data(request)
+    assert data_dict["team_id"] == team_id
+
+
+def test_prepare_mcp_server_data_update_excludes_none_team_id():
+    """UpdateMCPServerRequest without team_id should not have it in data dict."""
+    from litellm.proxy._experimental.mcp_server.db import _prepare_mcp_server_data
+    from litellm.proxy._types import UpdateMCPServerRequest
+
+    request = UpdateMCPServerRequest(
+        server_id="test-server-id",
+        transport="http",
+        url="https://example.com/mcp",
+    )
+    data_dict = _prepare_mcp_server_data(request)
+    assert "team_id" not in data_dict
+
+
+def test_update_request_model_fields_set_detects_explicit_null():
+    """When team_id is explicitly set to None in JSON, model_fields_set should contain it."""
+    from litellm.proxy._types import UpdateMCPServerRequest
+
+    # Simulate JSON: {"server_id": "x", "transport": "http", "url": "...", "team_id": null}
+    request = UpdateMCPServerRequest.model_validate(
+        {"server_id": "x", "transport": "http", "url": "https://example.com/mcp", "team_id": None}
+    )
+    assert "team_id" in request.model_fields_set
+
+    # Simulate JSON: {"server_id": "x", "transport": "http", "url": "..."} — team_id absent
+    request2 = UpdateMCPServerRequest.model_validate(
+        {"server_id": "x", "transport": "http", "url": "https://example.com/mcp"}
+    )
+    assert "team_id" not in request2.model_fields_set