Add pyright strict mode on the whole project (#1254)
This commit is contained in:
parent
ef4e167c1a
commit
c7671e470c
7 changed files with 652 additions and 674 deletions
|
|
@ -1,22 +1,22 @@
|
|||
from collections.abc import Callable, Generator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from anyio.streams.memory import MemoryObjectSendStream
|
||||
|
||||
import mcp.shared.memory
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.types import (
|
||||
JSONRPCNotification,
|
||||
JSONRPCRequest,
|
||||
)
|
||||
from mcp.types import JSONRPCNotification, JSONRPCRequest
|
||||
|
||||
|
||||
class SpyMemoryObjectSendStream:
|
||||
def __init__(self, original_stream):
|
||||
def __init__(self, original_stream: MemoryObjectSendStream[SessionMessage]):
|
||||
self.original_stream = original_stream
|
||||
self.sent_messages: list[SessionMessage] = []
|
||||
|
||||
async def send(self, message):
|
||||
async def send(self, message: SessionMessage):
|
||||
self.sent_messages.append(message)
|
||||
await self.original_stream.send(message)
|
||||
|
||||
|
|
@ -26,16 +26,12 @@ class SpyMemoryObjectSendStream:
|
|||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
async def __aexit__(self, *args: Any):
|
||||
await self.aclose()
|
||||
|
||||
|
||||
class StreamSpyCollection:
|
||||
def __init__(
|
||||
self,
|
||||
client_spy: SpyMemoryObjectSendStream,
|
||||
server_spy: SpyMemoryObjectSendStream,
|
||||
):
|
||||
def __init__(self, client_spy: SpyMemoryObjectSendStream, server_spy: SpyMemoryObjectSendStream):
|
||||
self.client = client_spy
|
||||
self.server = server_spy
|
||||
|
||||
|
|
@ -80,7 +76,7 @@ class StreamSpyCollection:
|
|||
|
||||
|
||||
@pytest.fixture
|
||||
def stream_spy():
|
||||
def stream_spy() -> Generator[Callable[[], StreamSpyCollection], None, None]:
|
||||
"""Fixture that provides spies for both client and server write streams.
|
||||
|
||||
Example usage:
|
||||
|
|
@ -103,7 +99,7 @@ def stream_spy():
|
|||
server_spy = None
|
||||
|
||||
# Store references to our spy objects
|
||||
def capture_spies(c_spy, s_spy):
|
||||
def capture_spies(c_spy: SpyMemoryObjectSendStream, s_spy: SpyMemoryObjectSendStream):
|
||||
nonlocal client_spy, server_spy
|
||||
client_spy = c_spy
|
||||
server_spy = s_spy
|
||||
|
|
|
|||
|
|
@ -11,12 +11,7 @@ from inline_snapshot import Is, snapshot
|
|||
from pydantic import AnyHttpUrl, AnyUrl
|
||||
|
||||
from mcp.client.auth import OAuthClientProvider, PKCEParameters
|
||||
from mcp.shared.auth import (
|
||||
OAuthClientInformationFull,
|
||||
OAuthClientMetadata,
|
||||
OAuthToken,
|
||||
ProtectedResourceMetadata,
|
||||
)
|
||||
from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken, ProtectedResourceMetadata
|
||||
|
||||
|
||||
class MockTokenStorage:
|
||||
|
|
@ -66,7 +61,7 @@ def valid_tokens():
|
|||
|
||||
|
||||
@pytest.fixture
|
||||
def oauth_provider(client_metadata, mock_storage):
|
||||
def oauth_provider(client_metadata: OAuthClientMetadata, mock_storage: MockTokenStorage):
|
||||
async def redirect_handler(url: str) -> None:
|
||||
"""Mock redirect handler."""
|
||||
pass
|
||||
|
|
@ -115,7 +110,9 @@ class TestOAuthContext:
|
|||
"""Test OAuth context functionality."""
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_oauth_provider_initialization(self, oauth_provider, client_metadata, mock_storage):
|
||||
async def test_oauth_provider_initialization(
|
||||
self, oauth_provider: OAuthClientProvider, client_metadata: OAuthClientMetadata, mock_storage: MockTokenStorage
|
||||
):
|
||||
"""Test OAuthClientProvider basic setup."""
|
||||
assert oauth_provider.context.server_url == "https://api.example.com/v1/mcp"
|
||||
assert oauth_provider.context.client_metadata == client_metadata
|
||||
|
|
@ -123,7 +120,7 @@ class TestOAuthContext:
|
|||
assert oauth_provider.context.timeout == 300.0
|
||||
assert oauth_provider.context is not None
|
||||
|
||||
def test_context_url_parsing(self, oauth_provider):
|
||||
def test_context_url_parsing(self, oauth_provider: OAuthClientProvider):
|
||||
"""Test get_authorization_base_url() extracts base URLs correctly."""
|
||||
context = oauth_provider.context
|
||||
|
||||
|
|
@ -145,7 +142,7 @@ class TestOAuthContext:
|
|||
)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_token_validity_checking(self, oauth_provider, mock_storage, valid_tokens):
|
||||
async def test_token_validity_checking(self, oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken):
|
||||
"""Test is_token_valid() and can_refresh_token() logic."""
|
||||
context = oauth_provider.context
|
||||
|
||||
|
|
@ -180,7 +177,7 @@ class TestOAuthContext:
|
|||
context.client_info = None
|
||||
assert not context.can_refresh_token()
|
||||
|
||||
def test_clear_tokens(self, oauth_provider, valid_tokens):
|
||||
def test_clear_tokens(self, oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken):
|
||||
"""Test clear_tokens() removes token data."""
|
||||
context = oauth_provider.context
|
||||
context.current_tokens = valid_tokens
|
||||
|
|
@ -198,7 +195,9 @@ class TestOAuthFlow:
|
|||
"""Test OAuth flow methods."""
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_discover_protected_resource_request(self, client_metadata, mock_storage):
|
||||
async def test_discover_protected_resource_request(
|
||||
self, client_metadata: OAuthClientMetadata, mock_storage: MockTokenStorage
|
||||
):
|
||||
"""Test protected resource discovery request building maintains backward compatibility."""
|
||||
|
||||
async def redirect_handler(url: str) -> None:
|
||||
|
|
@ -236,7 +235,7 @@ class TestOAuthFlow:
|
|||
assert "mcp-protocol-version" in request.headers
|
||||
|
||||
@pytest.mark.anyio
|
||||
def test_create_oauth_metadata_request(self, oauth_provider):
|
||||
def test_create_oauth_metadata_request(self, oauth_provider: OAuthClientProvider):
|
||||
"""Test OAuth metadata discovery request building."""
|
||||
request = oauth_provider._create_oauth_metadata_request("https://example.com")
|
||||
|
||||
|
|
@ -250,7 +249,7 @@ class TestOAuthFallback:
|
|||
"""Test OAuth discovery fallback behavior for legacy (act as AS not RS) servers."""
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_oauth_discovery_fallback_order(self, oauth_provider):
|
||||
async def test_oauth_discovery_fallback_order(self, oauth_provider: OAuthClientProvider):
|
||||
"""Test fallback URL construction order."""
|
||||
discovery_urls = oauth_provider._get_discovery_urls()
|
||||
|
||||
|
|
@ -262,7 +261,7 @@ class TestOAuthFallback:
|
|||
]
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_oauth_discovery_fallback_conditions(self, oauth_provider):
|
||||
async def test_oauth_discovery_fallback_conditions(self, oauth_provider: OAuthClientProvider):
|
||||
"""Test the conditions during which an AS metadata discovery fallback will be attempted."""
|
||||
# Ensure no tokens are stored
|
||||
oauth_provider.context.current_tokens = None
|
||||
|
|
@ -365,7 +364,7 @@ class TestOAuthFallback:
|
|||
token_request = await auth_flow.asend(token_response)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_handle_metadata_response_success(self, oauth_provider):
|
||||
async def test_handle_metadata_response_success(self, oauth_provider: OAuthClientProvider):
|
||||
"""Test successful metadata response handling."""
|
||||
# Create minimal valid OAuth metadata
|
||||
content = b"""{
|
||||
|
|
@ -381,7 +380,7 @@ class TestOAuthFallback:
|
|||
assert str(oauth_provider.context.oauth_metadata.issuer) == "https://auth.example.com/"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_register_client_request(self, oauth_provider):
|
||||
async def test_register_client_request(self, oauth_provider: OAuthClientProvider):
|
||||
"""Test client registration request building."""
|
||||
request = await oauth_provider._register_client()
|
||||
|
||||
|
|
@ -391,7 +390,7 @@ class TestOAuthFallback:
|
|||
assert request.headers["Content-Type"] == "application/json"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_register_client_skip_if_registered(self, oauth_provider, mock_storage):
|
||||
async def test_register_client_skip_if_registered(self, oauth_provider: OAuthClientProvider):
|
||||
"""Test client registration is skipped if already registered."""
|
||||
# Set existing client info
|
||||
client_info = OAuthClientInformationFull(
|
||||
|
|
@ -405,7 +404,7 @@ class TestOAuthFallback:
|
|||
assert request is None
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_token_exchange_request(self, oauth_provider):
|
||||
async def test_token_exchange_request(self, oauth_provider: OAuthClientProvider):
|
||||
"""Test token exchange request building."""
|
||||
# Set up required context
|
||||
oauth_provider.context.client_info = OAuthClientInformationFull(
|
||||
|
|
@ -429,7 +428,7 @@ class TestOAuthFallback:
|
|||
assert "client_secret=test_secret" in content
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_refresh_token_request(self, oauth_provider, valid_tokens):
|
||||
async def test_refresh_token_request(self, oauth_provider: OAuthClientProvider, valid_tokens: OAuthToken):
|
||||
"""Test refresh token request building."""
|
||||
# Set up required context
|
||||
oauth_provider.context.current_tokens = valid_tokens
|
||||
|
|
@ -538,11 +537,11 @@ class TestRegistrationResponse:
|
|||
"""Test client registration response handling."""
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_handle_registration_response_reads_before_accessing_text(self, oauth_provider):
|
||||
async def test_handle_registration_response_reads_before_accessing_text(self, oauth_provider: OAuthClientProvider):
|
||||
"""Test that response.aread() is called before accessing response.text."""
|
||||
|
||||
# Track if aread() was called
|
||||
class MockResponse:
|
||||
class MockResponse(httpx.Response):
|
||||
def __init__(self):
|
||||
self.status_code = 400
|
||||
self._aread_called = False
|
||||
|
|
@ -574,7 +573,9 @@ class TestAuthFlow:
|
|||
"""Test the auth flow in httpx."""
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_auth_flow_with_valid_tokens(self, oauth_provider, mock_storage, valid_tokens):
|
||||
async def test_auth_flow_with_valid_tokens(
|
||||
self, oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage, valid_tokens: OAuthToken
|
||||
):
|
||||
"""Test auth flow when tokens are already valid."""
|
||||
# Pre-store valid tokens
|
||||
await mock_storage.set_tokens(valid_tokens)
|
||||
|
|
@ -600,7 +601,7 @@ class TestAuthFlow:
|
|||
pass # Expected
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_auth_flow_with_no_tokens(self, oauth_provider, mock_storage):
|
||||
async def test_auth_flow_with_no_tokens(self, oauth_provider: OAuthClientProvider):
|
||||
"""Test auth flow when no tokens are available, triggering the full OAuth flow."""
|
||||
# Ensure no tokens are stored
|
||||
oauth_provider.context.current_tokens = None
|
||||
|
|
@ -810,7 +811,11 @@ class TestProtectedResourceWWWAuthenticate:
|
|||
],
|
||||
)
|
||||
def test_extract_resource_metadata_from_www_auth_valid_cases(
|
||||
self, client_metadata, mock_storage, www_auth_header, expected_url
|
||||
self,
|
||||
client_metadata: OAuthClientMetadata,
|
||||
mock_storage: MockTokenStorage,
|
||||
www_auth_header: str,
|
||||
expected_url: str,
|
||||
):
|
||||
"""Test extraction of resource_metadata URL from various valid WWW-Authenticate headers."""
|
||||
|
||||
|
|
@ -862,7 +867,12 @@ class TestProtectedResourceWWWAuthenticate:
|
|||
],
|
||||
)
|
||||
def test_extract_resource_metadata_from_www_auth_invalid_cases(
|
||||
self, client_metadata, mock_storage, status_code, www_auth_header, description
|
||||
self,
|
||||
client_metadata: OAuthClientMetadata,
|
||||
mock_storage: MockTokenStorage,
|
||||
status_code: int,
|
||||
www_auth_header: str | None,
|
||||
description: str,
|
||||
):
|
||||
"""Test extraction returns None for invalid cases."""
|
||||
|
||||
|
|
|
|||
|
|
@ -1,15 +1,16 @@
|
|||
from collections.abc import Callable
|
||||
|
||||
import pytest
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from mcp.shared.memory import (
|
||||
create_connected_server_and_client_session as create_session,
|
||||
)
|
||||
from mcp.shared.memory import create_connected_server_and_client_session as create_session
|
||||
|
||||
from .conftest import StreamSpyCollection
|
||||
|
||||
# Mark the whole module for async tests
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
async def test_list_tools_cursor_parameter(stream_spy):
|
||||
async def test_list_tools_cursor_parameter(stream_spy: Callable[[], StreamSpyCollection]):
|
||||
"""Test that the cursor parameter is accepted for list_tools
|
||||
and that it is correctly passed to the server.
|
||||
|
||||
|
|
@ -64,7 +65,7 @@ async def test_list_tools_cursor_parameter(stream_spy):
|
|||
assert list_tools_requests[0].params["cursor"] == ""
|
||||
|
||||
|
||||
async def test_list_resources_cursor_parameter(stream_spy):
|
||||
async def test_list_resources_cursor_parameter(stream_spy: Callable[[], StreamSpyCollection]):
|
||||
"""Test that the cursor parameter is accepted for list_resources
|
||||
and that it is correctly passed to the server.
|
||||
|
||||
|
|
@ -114,7 +115,7 @@ async def test_list_resources_cursor_parameter(stream_spy):
|
|||
assert list_resources_requests[0].params["cursor"] == ""
|
||||
|
||||
|
||||
async def test_list_prompts_cursor_parameter(stream_spy):
|
||||
async def test_list_prompts_cursor_parameter(stream_spy: Callable[[], StreamSpyCollection]):
|
||||
"""Test that the cursor parameter is accepted for list_prompts
|
||||
and that it is correctly passed to the server.
|
||||
See: https://modelcontextprotocol.io/specification/2025-03-26/server/utilities/pagination#request-format
|
||||
|
|
@ -163,7 +164,7 @@ async def test_list_prompts_cursor_parameter(stream_spy):
|
|||
assert list_prompts_requests[0].params["cursor"] == ""
|
||||
|
||||
|
||||
async def test_list_resource_templates_cursor_parameter(stream_spy):
|
||||
async def test_list_resource_templates_cursor_parameter(stream_spy: Callable[[], StreamSpyCollection]):
|
||||
"""Test that the cursor parameter is accepted for list_resource_templates
|
||||
and that it is correctly passed to the server.
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import logging
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -53,7 +54,7 @@ class TestClientOutputSchemaValidation:
|
|||
]
|
||||
|
||||
@server.call_tool()
|
||||
async def call_tool(name: str, arguments: dict):
|
||||
async def call_tool(name: str, arguments: dict[str, Any]):
|
||||
# Return invalid structured content - age is string instead of integer
|
||||
# The low-level server will wrap this in CallToolResult
|
||||
return {"name": "John", "age": "invalid"} # Invalid: age should be int
|
||||
|
|
@ -92,7 +93,7 @@ class TestClientOutputSchemaValidation:
|
|||
]
|
||||
|
||||
@server.call_tool()
|
||||
async def call_tool(name: str, arguments: dict):
|
||||
async def call_tool(name: str, arguments: dict[str, Any]):
|
||||
# Return invalid structured content - result is string instead of integer
|
||||
return {"result": "not_a_number"} # Invalid: should be int
|
||||
|
||||
|
|
@ -123,7 +124,7 @@ class TestClientOutputSchemaValidation:
|
|||
]
|
||||
|
||||
@server.call_tool()
|
||||
async def call_tool(name: str, arguments: dict):
|
||||
async def call_tool(name: str, arguments: dict[str, Any]):
|
||||
# Return invalid structured content - values should be integers
|
||||
return {"alice": "100", "bob": "85"} # Invalid: values should be int
|
||||
|
||||
|
|
@ -158,7 +159,7 @@ class TestClientOutputSchemaValidation:
|
|||
]
|
||||
|
||||
@server.call_tool()
|
||||
async def call_tool(name: str, arguments: dict):
|
||||
async def call_tool(name: str, arguments: dict[str, Any]):
|
||||
# Return structured content missing required field 'email'
|
||||
return {"name": "John", "age": 30} # Missing required 'email'
|
||||
|
||||
|
|
@ -170,17 +171,17 @@ class TestClientOutputSchemaValidation:
|
|||
assert "Invalid structured content returned by tool get_person" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_tool_not_listed_warning(self, caplog):
|
||||
async def test_tool_not_listed_warning(self, caplog: pytest.LogCaptureFixture):
|
||||
"""Test that client logs warning when tool is not in list_tools but has outputSchema"""
|
||||
server = Server("test-server")
|
||||
|
||||
@server.list_tools()
|
||||
async def list_tools():
|
||||
async def list_tools() -> list[Tool]:
|
||||
# Return empty list - tool is not listed
|
||||
return []
|
||||
|
||||
@server.call_tool()
|
||||
async def call_tool(name: str, arguments: dict):
|
||||
async def call_tool(name: str, arguments: dict[str, Any]) -> dict[str, Any]:
|
||||
# Server still responds to the tool call with structured content
|
||||
return {"result": 42}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,14 +1,12 @@
|
|||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
from mcp.shared.session import BaseSession
|
||||
from mcp.types import (
|
||||
ClientRequest,
|
||||
EmptyResult,
|
||||
PingRequest,
|
||||
)
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.shared.session import BaseSession, RequestId, SendResultT
|
||||
from mcp.types import ClientNotification, ClientRequest, ClientResult, EmptyResult, ErrorData, PingRequest
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
|
|
@ -20,13 +18,13 @@ async def test_send_request_stream_cleanup():
|
|||
"""
|
||||
|
||||
# Create a mock session with the minimal required functionality
|
||||
class TestSession(BaseSession):
|
||||
async def _send_response(self, request_id, response):
|
||||
class TestSession(BaseSession[ClientRequest, ClientNotification, ClientResult, Any, Any]):
|
||||
async def _send_response(self, request_id: RequestId, response: SendResultT | ErrorData) -> None:
|
||||
pass
|
||||
|
||||
# Create streams
|
||||
write_stream_send, write_stream_receive = anyio.create_memory_object_stream(1)
|
||||
read_stream_send, read_stream_receive = anyio.create_memory_object_stream(1)
|
||||
write_stream_send, write_stream_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
read_stream_send, read_stream_receive = anyio.create_memory_object_stream[SessionMessage](1)
|
||||
|
||||
# Create the session
|
||||
session = TestSession(
|
||||
|
|
@ -37,14 +35,10 @@ async def test_send_request_stream_cleanup():
|
|||
)
|
||||
|
||||
# Create a test request
|
||||
request = ClientRequest(
|
||||
PingRequest(
|
||||
method="ping",
|
||||
)
|
||||
)
|
||||
request = ClientRequest(PingRequest(method="ping"))
|
||||
|
||||
# Patch the _write_stream.send method to raise an exception
|
||||
async def mock_send(*args, **kwargs):
|
||||
async def mock_send(*args: Any, **kwargs: Any):
|
||||
raise RuntimeError("Simulated network error")
|
||||
|
||||
# Record the response streams before the test
|
||||
|
|
|
|||
|
|
@ -5,11 +5,7 @@ import pytest
|
|||
|
||||
import mcp
|
||||
from mcp import types
|
||||
from mcp.client.session_group import (
|
||||
ClientSessionGroup,
|
||||
SseServerParameters,
|
||||
StreamableHttpParameters,
|
||||
)
|
||||
from mcp.client.session_group import ClientSessionGroup, SseServerParameters, StreamableHttpParameters
|
||||
from mcp.client.stdio import StdioServerParameters
|
||||
from mcp.shared.exceptions import McpError
|
||||
|
||||
|
|
@ -54,7 +50,7 @@ class TestClientSessionGroup:
|
|||
mock_session = mock.AsyncMock()
|
||||
|
||||
# --- Prepare Session Group ---
|
||||
def hook(name, server_info):
|
||||
def hook(name: str, server_info: types.Implementation) -> str:
|
||||
return f"{(server_info.name)}-{name}"
|
||||
|
||||
mcp_session_group = ClientSessionGroup(component_name_hook=hook)
|
||||
|
|
@ -79,7 +75,7 @@ class TestClientSessionGroup:
|
|||
{"name": "value1", "args": {}},
|
||||
)
|
||||
|
||||
async def test_connect_to_server(self, mock_exit_stack):
|
||||
async def test_connect_to_server(self, mock_exit_stack: contextlib.AsyncExitStack):
|
||||
"""Test connecting to a server and aggregating components."""
|
||||
# --- Mock Dependencies ---
|
||||
mock_server_info = mock.Mock(spec=types.Implementation)
|
||||
|
|
@ -116,7 +112,7 @@ class TestClientSessionGroup:
|
|||
mock_session.list_resources.assert_awaited_once()
|
||||
mock_session.list_prompts.assert_awaited_once()
|
||||
|
||||
async def test_connect_to_server_with_name_hook(self, mock_exit_stack):
|
||||
async def test_connect_to_server_with_name_hook(self, mock_exit_stack: contextlib.AsyncExitStack):
|
||||
"""Test connecting with a component name hook."""
|
||||
# --- Mock Dependencies ---
|
||||
mock_server_info = mock.Mock(spec=types.Implementation)
|
||||
|
|
@ -208,7 +204,7 @@ class TestClientSessionGroup:
|
|||
assert "res1" not in group._resources
|
||||
assert "prm1" not in group._prompts
|
||||
|
||||
async def test_connect_to_server_duplicate_tool_raises_error(self, mock_exit_stack):
|
||||
async def test_connect_to_server_duplicate_tool_raises_error(self, mock_exit_stack: contextlib.AsyncExitStack):
|
||||
"""Test McpError raised when connecting a server with a dup name."""
|
||||
# --- Setup Pre-existing State ---
|
||||
group = ClientSessionGroup(exit_stack=mock_exit_stack)
|
||||
|
|
@ -282,9 +278,9 @@ class TestClientSessionGroup:
|
|||
)
|
||||
async def test_establish_session_parameterized(
|
||||
self,
|
||||
server_params_instance,
|
||||
client_type_name, # Just for clarity or conditional logic if needed
|
||||
patch_target_for_client_func,
|
||||
server_params_instance: StdioServerParameters | SseServerParameters | StreamableHttpParameters,
|
||||
client_type_name: str, # Just for clarity or conditional logic if needed
|
||||
patch_target_for_client_func: str,
|
||||
):
|
||||
with mock.patch("mcp.client.session_group.mcp.ClientSession") as mock_ClientSession_class:
|
||||
with mock.patch(patch_target_for_client_func) as mock_specific_client_func:
|
||||
|
|
@ -338,8 +334,10 @@ class TestClientSessionGroup:
|
|||
# --- Assertions ---
|
||||
# 1. Assert the correct specific client function was called
|
||||
if client_type_name == "stdio":
|
||||
assert isinstance(server_params_instance, StdioServerParameters)
|
||||
mock_specific_client_func.assert_called_once_with(server_params_instance)
|
||||
elif client_type_name == "sse":
|
||||
assert isinstance(server_params_instance, SseServerParameters)
|
||||
mock_specific_client_func.assert_called_once_with(
|
||||
url=server_params_instance.url,
|
||||
headers=server_params_instance.headers,
|
||||
|
|
@ -347,6 +345,7 @@ class TestClientSessionGroup:
|
|||
sse_read_timeout=server_params_instance.sse_read_timeout,
|
||||
)
|
||||
elif client_type_name == "streamablehttp":
|
||||
assert isinstance(server_params_instance, StreamableHttpParameters)
|
||||
mock_specific_client_func.assert_called_once_with(
|
||||
url=server_params_instance.url,
|
||||
headers=server_params_instance.headers,
|
||||
|
|
|
|||
|
|
@ -9,28 +9,25 @@ import anyio
|
|||
import pytest
|
||||
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.client.stdio import (
|
||||
StdioServerParameters,
|
||||
_create_platform_compatible_process,
|
||||
stdio_client,
|
||||
)
|
||||
from mcp.client.stdio import StdioServerParameters, _create_platform_compatible_process, stdio_client
|
||||
from mcp.shared.exceptions import McpError
|
||||
from mcp.shared.message import SessionMessage
|
||||
from mcp.types import CONNECTION_CLOSED, JSONRPCMessage, JSONRPCRequest, JSONRPCResponse
|
||||
from tests.shared.test_win32_utils import escape_path_for_python
|
||||
|
||||
from ..shared.test_win32_utils import escape_path_for_python
|
||||
|
||||
# Timeout for cleanup of processes that ignore SIGTERM
|
||||
# This timeout ensures the test fails quickly if the cleanup logic doesn't have
|
||||
# proper fallback mechanisms (SIGINT/SIGKILL) for processes that ignore SIGTERM
|
||||
SIGTERM_IGNORING_PROCESS_TIMEOUT = 5.0
|
||||
|
||||
tee: str = shutil.which("tee") # type: ignore
|
||||
python: str = shutil.which("python") # type: ignore
|
||||
tee = shutil.which("tee")
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.skipif(tee is None, reason="could not find tee command")
|
||||
async def test_stdio_context_manager_exiting():
|
||||
assert tee is not None
|
||||
async with stdio_client(StdioServerParameters(command=tee)) as (_, _):
|
||||
pass
|
||||
|
||||
|
|
@ -38,6 +35,7 @@ async def test_stdio_context_manager_exiting():
|
|||
@pytest.mark.anyio
|
||||
@pytest.mark.skipif(tee is None, reason="could not find tee command")
|
||||
async def test_stdio_client():
|
||||
assert tee is not None
|
||||
server_parameters = StdioServerParameters(command=tee)
|
||||
|
||||
async with stdio_client(server_parameters) as (read_stream, write_stream):
|
||||
|
|
@ -52,7 +50,7 @@ async def test_stdio_client():
|
|||
session_message = SessionMessage(message)
|
||||
await write_stream.send(session_message)
|
||||
|
||||
read_messages = []
|
||||
read_messages: list[JSONRPCMessage] = []
|
||||
async with read_stream:
|
||||
async for message in read_stream:
|
||||
if isinstance(message, Exception):
|
||||
|
|
@ -118,7 +116,7 @@ async def test_stdio_client_universal_cleanup():
|
|||
"""
|
||||
import time
|
||||
import sys
|
||||
|
||||
|
||||
# Simulate a long-running process
|
||||
for i in range(100):
|
||||
time.sleep(0.1)
|
||||
|
|
@ -136,7 +134,7 @@ async def test_stdio_client_universal_cleanup():
|
|||
start_time = time.time()
|
||||
|
||||
with anyio.move_on_after(8.0) as cancel_scope:
|
||||
async with stdio_client(server_params) as (read_stream, write_stream):
|
||||
async with stdio_client(server_params) as (_, _):
|
||||
# Immediately exit - this triggers cleanup while process is still running
|
||||
pass
|
||||
|
||||
|
|
@ -195,7 +193,7 @@ async def test_stdio_client_sigint_only_process():
|
|||
try:
|
||||
# Use anyio timeout to prevent test from hanging forever
|
||||
with anyio.move_on_after(5.0) as cancel_scope:
|
||||
async with stdio_client(server_params) as (read_stream, write_stream):
|
||||
async with stdio_client(server_params) as (_, _):
|
||||
# Let the process start and begin ignoring SIGTERM
|
||||
await anyio.sleep(0.5)
|
||||
# Exit context triggers cleanup - this should not hang
|
||||
|
|
@ -532,7 +530,7 @@ async def test_stdio_client_graceful_stdin_exit():
|
|||
script_content = textwrap.dedent(
|
||||
"""
|
||||
import sys
|
||||
|
||||
|
||||
# Read from stdin until it's closed
|
||||
try:
|
||||
while True:
|
||||
|
|
@ -541,7 +539,7 @@ async def test_stdio_client_graceful_stdin_exit():
|
|||
break
|
||||
except:
|
||||
pass
|
||||
|
||||
|
||||
# Exit gracefully
|
||||
sys.exit(0)
|
||||
"""
|
||||
|
|
@ -556,7 +554,7 @@ async def test_stdio_client_graceful_stdin_exit():
|
|||
|
||||
# Use anyio timeout to prevent test from hanging forever
|
||||
with anyio.move_on_after(5.0) as cancel_scope:
|
||||
async with stdio_client(server_params) as (read_stream, write_stream):
|
||||
async with stdio_client(server_params) as (_, _):
|
||||
# Let the process start and begin reading stdin
|
||||
await anyio.sleep(0.2)
|
||||
# Exit context triggers cleanup - process should exit from stdin closure
|
||||
|
|
@ -590,16 +588,16 @@ async def test_stdio_client_stdin_close_ignored():
|
|||
import signal
|
||||
import sys
|
||||
import time
|
||||
|
||||
|
||||
# Set up SIGTERM handler to exit cleanly
|
||||
def sigterm_handler(signum, frame):
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
signal.signal(signal.SIGTERM, sigterm_handler)
|
||||
|
||||
|
||||
# Close stdin immediately to simulate ignoring it
|
||||
sys.stdin.close()
|
||||
|
||||
|
||||
# Keep running until SIGTERM
|
||||
while True:
|
||||
time.sleep(0.1)
|
||||
|
|
@ -615,7 +613,7 @@ async def test_stdio_client_stdin_close_ignored():
|
|||
|
||||
# Use anyio timeout to prevent test from hanging forever
|
||||
with anyio.move_on_after(7.0) as cancel_scope:
|
||||
async with stdio_client(server_params) as (read_stream, write_stream):
|
||||
async with stdio_client(server_params) as (_, _):
|
||||
# Let the process start
|
||||
await anyio.sleep(0.2)
|
||||
# Exit context triggers cleanup
|
||||
|
|
|
|||
Loading…
Reference in a new issue