From 0fb6992f30f3925f83d90fc35f71e84b6c043f4c Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 26 Sep 2026 17:28:32 -0700 Subject: [PATCH] test(agents): use typed auth in MCP session fixtures --- .../mcp_server/test_mcp_server.py | 7 +++-- .../mcp_server/test_mcp_stale_session.py | 29 ++++++++++--------- 2 files changed, 19 insertions(+), 17 deletions(-) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index ab00ec4da1e..e7315c3c7da 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1,4 +1,5 @@ from litellm.proxy._experimental.mcp_server import operations as mcp_operations +from litellm.proxy._types import UserAPIKeyAuth import asyncio import contextlib import contextvars @@ -2244,7 +2245,7 @@ async def test_mcp_routing_chunked_initialize_to_stateful(): patch( "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", new_callable=AsyncMock, - return_value=(MagicMock(), None, ["progress_test"], None, None, None), + return_value=(UserAPIKeyAuth(), None, ["progress_test"], None, None, None), ), patch( "litellm.proxy._experimental.mcp_server.server.set_auth_context", @@ -2356,7 +2357,7 @@ async def test_mcp_routing_caps_body_peek_for_oversized_chunked_body(): patch( "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", new_callable=AsyncMock, - return_value=(MagicMock(), None, ["progress_test"], None, None, None), + return_value=(UserAPIKeyAuth(), None, ["progress_test"], None, None, None), ), patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"), patch( @@ -2567,7 +2568,7 @@ async def test_mcp_routing_initialize_rejected_when_owner_at_session_cap(): patch( "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", new_callable=AsyncMock, - return_value=(MagicMock(), None, ["progress_test"], None, None, None), + return_value=(UserAPIKeyAuth(), None, ["progress_test"], None, None, None), ), patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"), patch( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py index ec6fdef69ee..0124db1242a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_stale_session.py @@ -9,6 +9,7 @@ they may send a stale `mcp-session-id` header. This test verifies that: import asyncio from unittest.mock import AsyncMock, MagicMock, patch +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp import MCPAuth import pytest @@ -386,7 +387,7 @@ async def test_failed_delete_preserves_stateful_session_tracking(): pytest.skip("MCP server not available") session_id = "delete-failure-session" - user_auth = MagicMock() + user_auth = UserAPIKeyAuth() user_auth.api_key = "sk-test" user_auth.user_id = "test-user" auth_context = MagicMock() @@ -613,7 +614,7 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401(): } receive = AsyncMock() send = AsyncMock() - user_auth = MagicMock() + user_auth = UserAPIKeyAuth() user_auth.user_id = "test-user-id" oauth_server = MagicMock() oauth_server.auth_type = MCPAuth.oauth2 @@ -700,7 +701,7 @@ async def test_admitted_subject_missing_stored_token_challenged_with_resource_me } receive = AsyncMock() send = AsyncMock() - user_auth = MagicMock() + user_auth = UserAPIKeyAuth() user_auth.user_id = "sso-user-42" user_auth.mcp_admitted_user_subject = True oauth_server = MagicMock() @@ -806,7 +807,7 @@ async def test_client_credentials_server_is_not_preemptively_challenged(m2m_fiel } ) send = AsyncMock() - user_auth = MagicMock() + user_auth = UserAPIKeyAuth() user_auth.user_id = "test-user-id" m2m_server = MCPServer( server_id="m2m-server-id", @@ -892,7 +893,7 @@ async def test_handle_streamable_http_mcp_delegated_server_surfaces_upstream_cha } ) send = AsyncMock() - user_auth = MagicMock() + user_auth = UserAPIKeyAuth() user_auth.user_id = None delegated_server = MagicMock() delegated_server.auth_type = MCPAuth.oauth2 @@ -996,7 +997,7 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401(): } ) send = AsyncMock() - user_auth = MagicMock() + user_auth = UserAPIKeyAuth() user_auth.user_id = "test-user-id" oauth_server = MagicMock() oauth_server.auth_type = MCPAuth.oauth2 @@ -1092,7 +1093,7 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_returns } ) send = AsyncMock() - user_auth = MagicMock() + user_auth = UserAPIKeyAuth() user_auth.user_id = None delegated_server = MagicMock() delegated_server.auth_type = MCPAuth.oauth2 @@ -1192,7 +1193,7 @@ async def test_handle_streamable_http_mcp_token_exchange_without_subject_returns } ) send = AsyncMock() - user_auth = MagicMock() + user_auth = UserAPIKeyAuth() user_auth.user_id = None obo_server = MagicMock() obo_server.auth_type = MCPAuth.oauth2_token_exchange @@ -1301,7 +1302,7 @@ async def test_handle_streamable_http_mcp_oauth_delegate_without_token_returns_g } ) send = AsyncMock() - user_auth = MagicMock() + user_auth = UserAPIKeyAuth() user_auth.user_id = "u1" od_server = _build_passthrough_mode_server("od_server", MCPAuth.oauth_delegate) @@ -1366,7 +1367,7 @@ async def test_handle_streamable_http_mcp_oauth_delegate_with_forwarded_token_sk } ) send = AsyncMock() - user_auth = MagicMock() + user_auth = UserAPIKeyAuth() user_auth.user_id = "u1" od_server = _build_passthrough_mode_server("od_server", MCPAuth.oauth_delegate) @@ -1431,7 +1432,7 @@ async def _run_passthrough_connect( } ) send = AsyncMock() - user_auth = MagicMock() + user_auth = UserAPIKeyAuth() user_auth.user_id = "u1" server = _build_passthrough_mode_server(server_names[0], auth_type) @@ -1554,7 +1555,7 @@ async def test_handle_streamable_http_mcp_true_passthrough_without_token_surface } ) send = AsyncMock() - user_auth = MagicMock() + user_auth = UserAPIKeyAuth() user_auth.user_id = None tp_server = _build_passthrough_mode_server("tp_server", MCPAuth.true_passthrough) @@ -1620,7 +1621,7 @@ async def test_handle_streamable_http_mcp_true_passthrough_dcr_bridge_challenges } ) send = AsyncMock() - user_auth = MagicMock() + user_auth = UserAPIKeyAuth() user_auth.user_id = None bridge_server = _build_passthrough_mode_server("tp_bridge_server", MCPAuth.true_passthrough).model_copy( update={"dcr_bridge": True} @@ -1691,7 +1692,7 @@ async def test_handle_streamable_http_mcp_true_passthrough_with_token_skips_prob } ) send = AsyncMock() - user_auth = MagicMock() + user_auth = UserAPIKeyAuth() user_auth.user_id = None tp_server = _build_passthrough_mode_server("tp_server", MCPAuth.true_passthrough)