yihuang2026 commited on
Commit
a14756d
·
unverified ·
1 Parent(s): d4d4e69

add unit test

Browse files
src/fastmcp/client/client.py CHANGED
@@ -45,7 +45,7 @@ class Client:
45
  ):
46
  self.transport = infer_transport(transport)
47
  self._session: ClientSession | None = None
48
- self._session_cms: AbstractAsyncContextManager[ClientSession] | None = None
49
  self._nesting_counter: int = 0
50
 
51
  self._session_kwargs: SessionKwargs = {
@@ -95,11 +95,11 @@ class Client:
95
  return self
96
 
97
  async def __aexit__(self, exc_type, exc_val, exc_tb):
98
- self._nesting_counter -= 0
99
 
100
- if self._nesting_counter == 0 and self._session_cms is not None:
101
- await self._session_cms.__aexit__(exc_type, exc_val, exc_tb)
102
- self._session_cms = None
103
  self._session = None
104
 
105
  # --- MCP Client Methods ---
 
45
  ):
46
  self.transport = infer_transport(transport)
47
  self._session: ClientSession | None = None
48
+ self._session_cm: AbstractAsyncContextManager[ClientSession] | None = None
49
  self._nesting_counter: int = 0
50
 
51
  self._session_kwargs: SessionKwargs = {
 
95
  return self
96
 
97
  async def __aexit__(self, exc_type, exc_val, exc_tb):
98
+ self._nesting_counter -= 1
99
 
100
+ if self._nesting_counter == 0 and self._session_cm is not None:
101
+ await self._session_cm.__aexit__(exc_type, exc_val, exc_tb)
102
+ self._session_cm = None
103
  self._session = None
104
 
105
  # --- MCP Client Methods ---
tests/client/test_client.py CHANGED
@@ -1,10 +1,15 @@
1
  from typing import cast
 
 
 
 
 
2
 
3
  import pytest
4
  from pydantic import AnyUrl
5
 
6
  from fastmcp.client import Client
7
- from fastmcp.client.transports import FastMCPTransport
8
  from fastmcp.server.server import FastMCP
9
 
10
 
@@ -159,6 +164,38 @@ async def test_client_connection(fastmcp_server):
159
  # After connection
160
  assert not client.is_connected()
161
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
162
 
163
  async def test_resource_template(fastmcp_server):
164
  """Test using a resource template with InMemoryClient."""
 
1
  from typing import cast
2
+ from typing_extensions import Unpack
3
+ from collections.abc import AsyncIterator
4
+ from mcp import ClientSession
5
+ import contextlib
6
+ from mcp.shared.memory import create_client_server_memory_streams
7
 
8
  import pytest
9
  from pydantic import AnyUrl
10
 
11
  from fastmcp.client import Client
12
+ from fastmcp.client.transports import FastMCPTransport, ClientTransport, SessionKwargs
13
  from fastmcp.server.server import FastMCP
14
 
15
 
 
164
  # After connection
165
  assert not client.is_connected()
166
 
167
+ async def test_client_nested_context_manager(fastmcp_server):
168
+ """Test that the client connects and disconnects once in nested context manager."""
169
+ class MockTransport(ClientTransport):
170
+ def __init__(self):
171
+ self._connected = False
172
+
173
+ @contextlib.asynccontextmanager
174
+ async def connect_session(
175
+ self, **session_kwargs: Unpack[SessionKwargs],
176
+ ) -> AsyncIterator[ClientSession]:
177
+ assert not self._connected, "Transport is connected multiple times"
178
+ self._connected = True
179
+ async with create_client_server_memory_streams() as (
180
+ _,
181
+ server_streams,
182
+ ):
183
+ yield ClientSession(*server_streams)
184
+
185
+ client = Client(transport=MockTransport())
186
+
187
+ # Before connection
188
+ assert not client.is_connected()
189
+
190
+ # During connection
191
+ async with client:
192
+ assert client.is_connected()
193
+
194
+ async with client:
195
+ assert client.is_connected()
196
+
197
+ # After connection
198
+ assert not client.is_connected()
199
 
200
  async def test_resource_template(fastmcp_server):
201
  """Test using a resource template with InMemoryClient."""