diff --git a/tests/client/test_client.py b/tests/client/test_client.py index 5165c2a5d..7bc0fe22e 100644 --- a/tests/client/test_client.py +++ b/tests/client/test_client.py @@ -1,15 +1,10 @@ -import contextlib -from collections.abc import AsyncIterator from typing import cast import pytest -from mcp import ClientSession -from mcp.shared.memory import create_client_server_memory_streams from pydantic import AnyUrl -from typing_extensions import Unpack from fastmcp.client import Client -from fastmcp.client.transports import ClientTransport, FastMCPTransport, SessionKwargs +from fastmcp.client.transports import FastMCPTransport from fastmcp.server.server import FastMCP @@ -168,40 +163,31 @@ async def test_client_connection(fastmcp_server): async def test_client_nested_context_manager(fastmcp_server): """Test that the client connects and disconnects once in nested context manager.""" - class MockTransport(ClientTransport): - def __init__(self): - self._connected = False - - @contextlib.asynccontextmanager - async def connect_session( - self, - **session_kwargs: Unpack[SessionKwargs], - ) -> AsyncIterator[ClientSession]: - assert not self._connected, "Transport is connected multiple times" - self._connected = True - async with create_client_server_memory_streams() as ( - _, - server_streams, - ): - yield ClientSession(*server_streams) - - client = Client(transport=MockTransport()) + client = Client(fastmcp_server) # Before connection assert not client.is_connected() + assert client._session is None # During connection async with client: assert client.is_connected() + assert client._session is not None + session = client._session + # Re-use the same session async with client: assert client.is_connected() + assert client._session is session + # Re-use the same session async with client: assert client.is_connected() + assert client._session is session # After connection assert not client.is_connected() + assert client._session is None async def test_resource_template(fastmcp_server):