From e2605c02bb1e1bf701544631817cc3bb6805f1a0 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 11 Apr 2026 11:40:44 -0700 Subject: [PATCH] fix(tests/mcp): mock find_unique to return existing server record in token exchange tests --- .../mcp_server/test_discoverable_endpoints.py | 56 +++++++++++++------ 1 file changed, 40 insertions(+), 16 deletions(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 4ff57b0c3f6..b76192555ea 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -10,7 +10,7 @@ from fastapi import HTTPException @pytest.fixture(autouse=True) def mock_mcp_client_ip(): """Mock IPAddressUtils.get_mcp_client_ip to return None for all tests. - + This bypasses IP-based access control in tests, since the MCP server's available_on_public_internet defaults to False and mock requests don't have proper client IP context. @@ -145,9 +145,9 @@ async def test_authorize_endpoint_preserves_existing_query_params(): location = response.headers["location"] # Must NOT have double '?' — existing params must be merged correctly - assert location.count("?") == 1, ( - f"Expected exactly one '?' in URL but got {location.count('?')}: {location}" - ) + assert ( + location.count("?") == 1 + ), f"Expected exactly one '?' in URL but got {location.count('?')}: {location}" assert "tenant=system" in location assert "client_id=test_client_id" in location assert "response_type=code" in location @@ -1521,7 +1521,10 @@ async def test_oauth_callback_redirects_with_state(): # Should redirect to the client callback URL with code and original state assert response.status_code == 302 - assert "http://localhost:3000/ui/mcp/oauth/callback" in response.headers["location"] + assert ( + "http://localhost:3000/ui/mcp/oauth/callback" + in response.headers["location"] + ) assert "code=test_authorization_code_12345" in response.headers["location"] assert "state=test-uuid-state-123" in response.headers["location"] @@ -1609,7 +1612,10 @@ async def test_oauth_authorize_includes_scopes_from_server_config(): # Should redirect with scopes from server config assert response.status_code in (307, 302) redirect_url = response.headers["location"] - assert "scope=api+read_user+ai_workflows" in redirect_url or "scope=api%20read_user%20ai_workflows" in redirect_url + assert ( + "scope=api+read_user+ai_workflows" in redirect_url + or "scope=api%20read_user%20ai_workflows" in redirect_url + ) @pytest.mark.asyncio @@ -1664,7 +1670,10 @@ async def test_oauth_authorize_prefers_request_scope_over_server_config(): # Should use the explicit scope, not server config assert response.status_code in (307, 302) redirect_url = response.headers["location"] - assert "scope=custom_scope1+custom_scope2" in redirect_url or "scope=custom_scope1%20custom_scope2" in redirect_url + assert ( + "scope=custom_scope1+custom_scope2" in redirect_url + or "scope=custom_scope1%20custom_scope2" in redirect_url + ) assert "default_scope" not in redirect_url @@ -1849,10 +1858,17 @@ async def test_token_exchange_persists_credentials_to_db(): } mock_response.raise_for_status = MagicMock() + mock_server_record = MagicMock() + mock_server_record.server_id = "test-persist-server" + mock_server_record.credentials = ( + None # No existing credentials — first token exchange + ) + mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=None) - # No existing record — first token exchange - mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock( + return_value=mock_server_record + ) with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" @@ -1950,9 +1966,15 @@ async def test_token_exchange_persists_without_optional_fields(): } mock_response.raise_for_status = MagicMock() + mock_server_record = MagicMock() + mock_server_record.server_id = "test-persist-minimal" + mock_server_record.credentials = None + mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=None) - mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock( + return_value=mock_server_record + ) with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" @@ -2035,12 +2057,14 @@ async def test_token_refresh_preserves_existing_refresh_token(): # Existing DB record has the original refresh_token mock_existing = MagicMock() - mock_existing.credentials = json.dumps({ - "auth_value": "old_encrypted_access", - "refresh_token": "original_refresh_token_encrypted", - "client_id": "existing_client_id", - "type": "oauth2", - }) + mock_existing.credentials = json.dumps( + { + "auth_value": "old_encrypted_access", + "refresh_token": "original_refresh_token_encrypted", + "client_id": "existing_client_id", + "type": "oauth2", + } + ) mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=None)