mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-22 05:24:18 +02:00
add unit test
This commit is contained in:
parent
0af58a85a6
commit
5bbad74044
2 changed files with 43 additions and 6 deletions
|
|
@ -45,7 +45,7 @@ class Client:
|
|||
):
|
||||
self.transport = infer_transport(transport)
|
||||
self._session: ClientSession | None = None
|
||||
self._session_cms: AbstractAsyncContextManager[ClientSession] | None = None
|
||||
self._session_cm: AbstractAsyncContextManager[ClientSession] | None = None
|
||||
self._nesting_counter: int = 0
|
||||
|
||||
self._session_kwargs: SessionKwargs = {
|
||||
|
|
@ -95,11 +95,11 @@ class Client:
|
|||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
self._nesting_counter -= 0
|
||||
self._nesting_counter -= 1
|
||||
|
||||
if self._nesting_counter == 0 and self._session_cms is not None:
|
||||
await self._session_cms.__aexit__(exc_type, exc_val, exc_tb)
|
||||
self._session_cms = None
|
||||
if self._nesting_counter == 0 and self._session_cm is not None:
|
||||
await self._session_cm.__aexit__(exc_type, exc_val, exc_tb)
|
||||
self._session_cm = None
|
||||
self._session = None
|
||||
|
||||
# --- MCP Client Methods ---
|
||||
|
|
|
|||
|
|
@ -1,10 +1,15 @@
|
|||
from typing import cast
|
||||
from typing_extensions import Unpack
|
||||
from collections.abc import AsyncIterator
|
||||
from mcp import ClientSession
|
||||
import contextlib
|
||||
from mcp.shared.memory import create_client_server_memory_streams
|
||||
|
||||
import pytest
|
||||
from pydantic import AnyUrl
|
||||
|
||||
from fastmcp.client import Client
|
||||
from fastmcp.client.transports import FastMCPTransport
|
||||
from fastmcp.client.transports import FastMCPTransport, ClientTransport, SessionKwargs
|
||||
from fastmcp.server.server import FastMCP
|
||||
|
||||
|
||||
|
|
@ -159,6 +164,38 @@ async def test_client_connection(fastmcp_server):
|
|||
# After connection
|
||||
assert not client.is_connected()
|
||||
|
||||
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())
|
||||
|
||||
# Before connection
|
||||
assert not client.is_connected()
|
||||
|
||||
# During connection
|
||||
async with client:
|
||||
assert client.is_connected()
|
||||
|
||||
async with client:
|
||||
assert client.is_connected()
|
||||
|
||||
# After connection
|
||||
assert not client.is_connected()
|
||||
|
||||
async def test_resource_template(fastmcp_server):
|
||||
"""Test using a resource template with InMemoryClient."""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue