|
5 | 5 | import multiprocessing |
6 | 6 | import socket |
7 | 7 | import time |
8 | | -from collections.abc import Callable, Generator |
9 | | -from contextlib import contextmanager |
| 8 | +from collections.abc import AsyncGenerator, Callable, Generator |
| 9 | +from contextlib import asynccontextmanager, contextmanager |
10 | 10 | from typing import TYPE_CHECKING, Any, Literal |
11 | 11 | from urllib.parse import parse_qs, urlparse |
12 | 12 |
|
@@ -140,6 +140,88 @@ def run_server_in_process( |
140 | 140 | raise RuntimeError("Server process failed to terminate even after kill") |
141 | 141 |
|
142 | 142 |
|
| 143 | +@asynccontextmanager |
| 144 | +async def run_server_async( |
| 145 | + server: FastMCP, |
| 146 | + port: int | None = None, |
| 147 | + transport: Literal["http", "streamable-http", "sse"] = "http", |
| 148 | + path: str = "/mcp", |
| 149 | + host: str = "127.0.0.1", |
| 150 | +) -> AsyncGenerator[str, None]: |
| 151 | + """ |
| 152 | + Start a FastMCP server as an asyncio task for in-process async testing. |
| 153 | +
|
| 154 | + This is the recommended way to test FastMCP servers. It runs the server |
| 155 | + as an async task in the same process, eliminating subprocess coordination, |
| 156 | + sleeps, and cleanup issues. |
| 157 | +
|
| 158 | + Args: |
| 159 | + server: FastMCP server instance |
| 160 | + port: Port to bind to (default: find available port) |
| 161 | + transport: Transport type ("http", "streamable-http", or "sse") |
| 162 | + path: URL path for the server (default: "/mcp") |
| 163 | + host: Host to bind to (default: "127.0.0.1") |
| 164 | +
|
| 165 | + Yields: |
| 166 | + Server URL string |
| 167 | +
|
| 168 | + Example: |
| 169 | + ```python |
| 170 | + import pytest |
| 171 | + from fastmcp import FastMCP, Client |
| 172 | + from fastmcp.client.transports import StreamableHttpTransport |
| 173 | + from fastmcp.utilities.tests import run_server_async |
| 174 | +
|
| 175 | + @pytest.fixture |
| 176 | + async def server(): |
| 177 | + mcp = FastMCP("test") |
| 178 | +
|
| 179 | + @mcp.tool() |
| 180 | + def greet(name: str) -> str: |
| 181 | + return f"Hello, {name}!" |
| 182 | +
|
| 183 | + async with run_server_async(mcp) as url: |
| 184 | + yield url |
| 185 | +
|
| 186 | + async def test_greet(server: str): |
| 187 | + async with Client(StreamableHttpTransport(server)) as client: |
| 188 | + result = await client.call_tool("greet", {"name": "World"}) |
| 189 | + assert result.content[0].text == "Hello, World!" |
| 190 | + ``` |
| 191 | + """ |
| 192 | + import asyncio |
| 193 | + |
| 194 | + if port is None: |
| 195 | + port = find_available_port() |
| 196 | + |
| 197 | + # Wait a tiny bit for the port to be released if it was just used |
| 198 | + await asyncio.sleep(0.01) |
| 199 | + |
| 200 | + # Start server as a background task |
| 201 | + server_task = asyncio.create_task( |
| 202 | + server.run_http_async( |
| 203 | + host=host, |
| 204 | + port=port, |
| 205 | + transport=transport, |
| 206 | + path=path, |
| 207 | + show_banner=False, |
| 208 | + ) |
| 209 | + ) |
| 210 | + |
| 211 | + # Give the server a moment to start |
| 212 | + await asyncio.sleep(0.1) |
| 213 | + |
| 214 | + try: |
| 215 | + yield f"http://{host}:{port}{path}" |
| 216 | + finally: |
| 217 | + # Cleanup: cancel the task |
| 218 | + server_task.cancel() |
| 219 | + try: |
| 220 | + await server_task |
| 221 | + except asyncio.CancelledError: |
| 222 | + pass |
| 223 | + |
| 224 | + |
143 | 225 | @contextmanager |
144 | 226 | def caplog_for_fastmcp(caplog): |
145 | 227 | """Context manager to capture logs from FastMCP loggers even when propagation is disabled.""" |
|
0 commit comments