Spaces:
Running
Running
add unit test
Browse files- src/fastmcp/client/client.py +5 -5
- tests/client/test_client.py +38 -1
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.
|
| 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 -=
|
| 99 |
|
| 100 |
-
if self._nesting_counter == 0 and self.
|
| 101 |
-
await self.
|
| 102 |
-
self.
|
| 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."""
|