diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 2eab053ab73..5fab4ceec72 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py @@ -19,6 +19,7 @@ from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.outbound_credentials import ( ApiKeyConfig, + AuthConfig, AuthorizationCodeConfig, AwsSigV4Config, Byok, @@ -1216,7 +1217,9 @@ async def test_passthrough_ignores_the_carrier_and_keeps_the_callers_slot(): (AuthorizationCodeConfig(), Subject(tenant_id="", subject_id="alice"), "stored-user-token", "Bearer stored-alice"), ], ) -async def test_resolved_source_matches_the_credential_sent_upstream(config, subject, expected_source, expected_header): +async def test_resolved_source_matches_the_credential_sent_upstream( + config: AuthConfig, subject: Subject, expected_source: str, expected_header: str | None +) -> None: from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source store = _FakeTokenStore({("alice", "s"): OAuthToken(access_token="stored-alice")}) @@ -1230,7 +1233,7 @@ async def test_resolved_source_matches_the_credential_sent_upstream(config, subj @pytest.mark.asyncio -async def test_resolved_source_preserves_missing_user_token_error(): +async def test_resolved_source_preserves_missing_user_token_error() -> None: from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source result = await resolve_credentials_with_source(UpstreamCredentialProvider(), _SUBJECT, _spec(AuthorizationCodeConfig())) @@ -1239,7 +1242,7 @@ async def test_resolved_source_preserves_missing_user_token_error(): @pytest.mark.asyncio -async def test_minted_token_sources_match_egress_and_do_not_fetch_twice(): +async def test_minted_token_sources_match_egress_and_do_not_fetch_twice() -> None: from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source source = _FakeM2MSource(Ok(OAuthToken(access_token="m2m-at"))) @@ -1259,7 +1262,7 @@ async def test_minted_token_sources_match_egress_and_do_not_fetch_twice(): @pytest.mark.asyncio -async def test_id_jag_source_describes_final_token_after_both_exchanges(): +async def test_id_jag_source_describes_final_token_after_both_exchanges() -> None: from litellm.proxy._experimental.mcp_server.outbound_credentials.resolver import resolve_credentials_with_source endpoint = _FakeTokenEndpoint(_two_leg_ok("resource-token")) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 99447acd371..adf985e9a21 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -5,7 +5,7 @@ import logging import os import sys from datetime import datetime -from typing import Any, Dict, Final, Optional +from typing import Any, Dict, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -12582,8 +12582,11 @@ async def test_pre_call_tool_check_honors_guardrail_attached_to_key(monkeypatch, ], ) async def test_debug_resolution_matches_final_header_conflict_winner( - config, extra_headers, expected_source, expected_authorization -): + config: Literal["stored", "static", "none"], + extra_headers: dict[str, str] | None, + expected_source: str, + expected_authorization: str | None, +) -> None: from mcp.server.lowlevel.server import request_ctx from mcp.shared.context import RequestContext from starlette.requests import Request @@ -12598,7 +12601,7 @@ async def test_debug_resolution_matches_final_header_conflict_winner( from litellm.types.mcp_server.mcp_server_manager import MCPServer class Store: - def __init__(self): + def __init__(self) -> None: self.calls = 0 async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: @@ -12641,7 +12644,7 @@ async def test_debug_resolution_matches_final_header_conflict_winner( @pytest.mark.asyncio @pytest.mark.parametrize("transport", ["http", "stdio"]) -async def test_debug_reports_legacy_signing_and_non_http_transport(transport): +async def test_debug_reports_legacy_signing_and_non_http_transport(transport: Literal["http", "stdio"]) -> None: from mcp.server.lowlevel.server import request_ctx from mcp.shared.context import RequestContext from starlette.requests import Request