From 8fb709c98c8162d5a0a01f21dedddfc21b375723 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sun, 1 Jun 2025 09:35:06 -0400 Subject: [PATCH] Simplify code for running servers in processes --- src/fastmcp/utilities/tests.py | 22 ++++-- tests/auth/test_oauth_client.py | 25 +------ tests/client/test_openapi.py | 77 +++++---------------- tests/client/test_sse.py | 36 +++------- tests/client/test_streamable_http.py | 56 +++++---------- tests/server/http/test_http_dependencies.py | 44 ++---------- 6 files changed, 70 insertions(+), 190 deletions(-) diff --git a/src/fastmcp/utilities/tests.py b/src/fastmcp/utilities/tests.py index 4fa73006e..149182f3b 100644 --- a/src/fastmcp/utilities/tests.py +++ b/src/fastmcp/utilities/tests.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, Literal import uvicorn +import fastmcp from fastmcp.settings import settings from fastmcp.utilities.http import find_available_port @@ -72,14 +73,19 @@ def _run_server(mcp_server: FastMCP, transport: Literal["sse"], port: int) -> No @contextmanager def run_server_in_process( - server_fn: Callable[..., None], *args + server_fn: FastMCP | Callable[..., None], + *args, + provide_host_and_port: bool = True, + **kwargs, ) -> Generator[str, None, None]: """ - Context manager that runs a Starlette app in a separate process and returns the - server URL. When the context manager is exited, the server process is killed. + Context manager that runs a FastMCP server (or a function that runs a FastMCP server) in a separate process and + returns the server URL. When the context manager is exited, the server process is killed. Args: - app: The Starlette app to run. + server_fn: The FastMCP server to run, or a function that runs a FastMCP + server. If a FastMCP server is provided, its .run() method is called + with the provided arguments and keyword arguments. Returns: The server URL. @@ -87,8 +93,14 @@ def run_server_in_process( host = "127.0.0.1" port = find_available_port() + if isinstance(server_fn, fastmcp.FastMCP): + server_fn = server_fn.run + + if provide_host_and_port: + kwargs |= {"host": host, "port": port} + proc = multiprocessing.Process( - target=server_fn, args=(host, port, *args), daemon=True + target=server_fn, args=args, kwargs=kwargs, daemon=True ) proc.start() diff --git a/tests/auth/test_oauth_client.py b/tests/auth/test_oauth_client.py index 2a1fd4c23..5d668c6dc 100644 --- a/tests/auth/test_oauth_client.py +++ b/tests/auth/test_oauth_client.py @@ -1,11 +1,9 @@ -import sys from collections.abc import Generator from unittest.mock import patch from urllib.parse import parse_qs, urlparse import httpx import pytest -import uvicorn import fastmcp.client.auth # Import module, not the function directly from fastmcp.client import Client @@ -39,30 +37,13 @@ def fastmcp_server(issuer_url: str): return server -def run_server(host: str, port: int, transport: str | None = None) -> None: - try: - # Configure OAuth provider with the actual server URL - issuer_url = f"http://{host}:{port}" - app = fastmcp_server(issuer_url).http_app() - server = uvicorn.Server( - config=uvicorn.Config( - app=app, - host=host, - port=port, - log_level="error", - lifespan="on", - ) - ) - server.run() - except Exception as e: - print(f"Server error: {e}") - sys.exit(1) - sys.exit(0) +def run_server(host: str, port: int, **kwargs) -> None: + fastmcp_server(f"http://{host}:{port}").run(host=host, port=port, **kwargs) @pytest.fixture(scope="module") def streamable_http_server() -> Generator[str, None, None]: - with run_server_in_process(run_server) as url: + with run_server_in_process(run_server, transport="streamable-http") as url: yield f"{url}/mcp" diff --git a/tests/client/test_openapi.py b/tests/client/test_openapi.py index 8ea642096..000895e87 100644 --- a/tests/client/test_openapi.py +++ b/tests/client/test_openapi.py @@ -1,9 +1,7 @@ import json -import sys from collections.abc import Generator import pytest -import uvicorn from fastapi import FastAPI, Request from fastmcp import Client, FastMCP @@ -34,75 +32,34 @@ def fastmcp_server_for_headers() -> FastMCP: return mcp +def run_server(host: str, port: int, **kwargs) -> None: + fastmcp_server_for_headers().run(host=host, port=port, **kwargs) + + +def run_proxy_server(host: str, port: int, shttp_url: str, **kwargs) -> None: + client = Client(transport=StreamableHttpTransport(shttp_url)) + app = FastMCP.as_proxy(client) + app.run(host=host, port=port, **kwargs) + + class TestClientHeaders: - def run_shttp_server(self, host: str, port: int) -> None: - try: - app = fastmcp_server_for_headers().http_app(transport="streamable-http") - server = uvicorn.Server( - config=uvicorn.Config( - app=app, - host=host, - port=port, - log_level="error", - lifespan="on", - ) - ) - server.run() - except Exception as e: - print(f"Server error: {e}") - sys.exit(1) - sys.exit(0) - - def run_sse_server(self, host: str, port: int) -> None: - try: - app = fastmcp_server_for_headers().http_app(transport="sse") - server = uvicorn.Server( - config=uvicorn.Config( - app=app, - host=host, - port=port, - log_level="error", - lifespan="on", - ) - ) - server.run() - except Exception as e: - print(f"Server error: {e}") - sys.exit(1) - sys.exit(0) - - def run_proxy_server(self, host: str, port: int, remote_url: str) -> None: - try: - client = Client(transport=StreamableHttpTransport(remote_url)) - app = FastMCP.as_proxy(client).http_app(transport="streamable-http") - server = uvicorn.Server( - config=uvicorn.Config( - app=app, - host=host, - port=port, - log_level="error", - lifespan="on", - ) - ) - server.run() - except Exception as e: - print(f"Server error: {e}") - sys.exit(1) - sys.exit(0) - @pytest.fixture(scope="class") def shttp_server(self) -> Generator[str, None, None]: - with run_server_in_process(self.run_shttp_server) as url: + with run_server_in_process(run_server, transport="streamable-http") as url: yield f"{url}/mcp" @pytest.fixture(scope="class") def sse_server(self) -> Generator[str, None, None]: - with run_server_in_process(self.run_sse_server) as url: + with run_server_in_process(run_server, transport="sse") as url: yield f"{url}/sse" @pytest.fixture(scope="class") def proxy_server(self, shttp_server: str) -> Generator[str, None, None]: - with run_server_in_process(self.run_proxy_server, shttp_server + "/mcp") as url: + with run_server_in_process( + run_proxy_server, + shttp_url=shttp_server, + transport="streamable-http", + ) as url: yield f"{url}/mcp" async def test_client_headers_sse_resource(self, sse_server: str): diff --git a/tests/client/test_sse.py b/tests/client/test_sse.py index ebf7a0a31..39b556dda 100644 --- a/tests/client/test_sse.py +++ b/tests/client/test_sse.py @@ -63,22 +63,13 @@ def fastmcp_server(): return server -def run_server(host: str, port: int, path: str | None = None) -> None: - try: - app = fastmcp_server().http_app(transport="sse", path=path) - server = uvicorn.Server( - config=uvicorn.Config(app=app, host=host, port=port, log_level="error") - ) - server.run() - except Exception as e: - print(f"Server error: {e}") - sys.exit(1) - sys.exit(0) +def run_server(host: str, port: int, **kwargs) -> None: + fastmcp_server().run(host=host, port=port, **kwargs) @pytest.fixture(autouse=True, scope="module") def sse_server() -> Generator[str, None, None]: - with run_server_in_process(run_server) as url: + with run_server_in_process(run_server, transport="sse") as url: yield f"{url}/sse" @@ -101,22 +92,17 @@ async def test_http_headers(sse_server: str): def run_nested_server(host: str, port: int) -> None: - try: - app = fastmcp_server().sse_app(path="/mcp/sse", message_path="/mcp/messages") - mount = Starlette(routes=[Mount("/nest-inner", app=app)]) - mount2 = Starlette(routes=[Mount("/nest-outer", app=mount)]) - server = uvicorn.Server( - config=uvicorn.Config(app=mount2, host=host, port=port, log_level="error") - ) - server.run() - except Exception as e: - print(f"Server error: {e}") - sys.exit(1) - sys.exit(0) + app = fastmcp_server().sse_app(path="/mcp/sse", message_path="/mcp/messages") + mount = Starlette(routes=[Mount("/nest-inner", app=app)]) + mount2 = Starlette(routes=[Mount("/nest-outer", app=mount)]) + server = uvicorn.Server( + config=uvicorn.Config(app=mount2, host=host, port=port, log_level="error") + ) + server.run() async def test_run_server_on_path(): - with run_server_in_process(run_server, "/help") as url: + with run_server_in_process(run_server, transport="sse", path="/help") as url: async with Client(transport=SSETransport(f"{url}/help")) as client: result = await client.ping() assert result is True diff --git a/tests/client/test_streamable_http.py b/tests/client/test_streamable_http.py index 34723e7c4..2c4e349e3 100644 --- a/tests/client/test_streamable_http.py +++ b/tests/client/test_streamable_http.py @@ -63,28 +63,13 @@ def fastmcp_server(): return server -def run_server(host: str, port: int) -> None: - try: - app = fastmcp_server().http_app() - server = uvicorn.Server( - config=uvicorn.Config( - app=app, - host=host, - port=port, - log_level="error", - lifespan="on", - ) - ) - server.run() - except Exception as e: - print(f"Server error: {e}") - sys.exit(1) - sys.exit(0) +def run_server(host: str, port: int, **kwargs) -> None: + fastmcp_server().run(host=host, port=port, **kwargs) @pytest.fixture(scope="module") def streamable_http_server() -> Generator[str, None, None]: - with run_server_in_process(run_server) as url: + with run_server_in_process(run_server, transport="streamable-http") as url: yield f"{url}/mcp" @@ -111,28 +96,23 @@ async def test_http_headers(streamable_http_server: str): def run_nested_server(host: str, port: int) -> None: - try: - mcp_app = fastmcp_server().http_app(path="/final/mcp") + mcp_app = fastmcp_server().http_app(path="/final/mcp") - mount = Starlette(routes=[Mount("/nest-inner", app=mcp_app)]) - mount2 = Starlette( - routes=[Mount("/nest-outer", app=mount)], - lifespan=mcp_app.lifespan, + mount = Starlette(routes=[Mount("/nest-inner", app=mcp_app)]) + mount2 = Starlette( + routes=[Mount("/nest-outer", app=mount)], + lifespan=mcp_app.lifespan, + ) + server = uvicorn.Server( + config=uvicorn.Config( + app=mount2, + host=host, + port=port, + log_level="error", + lifespan="on", ) - server = uvicorn.Server( - config=uvicorn.Config( - app=mount2, - host=host, - port=port, - log_level="error", - lifespan="on", - ) - ) - server.run() - except Exception as e: - print(f"Server error: {e}") - sys.exit(1) - sys.exit(0) + ) + server.run() async def test_nested_streamable_http_server_resolves_correctly(): diff --git a/tests/server/http/test_http_dependencies.py b/tests/server/http/test_http_dependencies.py index 938b55c14..580ceabd3 100644 --- a/tests/server/http/test_http_dependencies.py +++ b/tests/server/http/test_http_dependencies.py @@ -1,9 +1,7 @@ import json -import sys from collections.abc import Generator import pytest -import uvicorn from fastmcp.client import Client from fastmcp.client.transports import SSETransport, StreamableHttpTransport @@ -40,53 +38,19 @@ def fastmcp_server(): return server -def run_shttp_server(host: str, port: int) -> None: - try: - app = fastmcp_server().http_app(transport="streamable-http") - server = uvicorn.Server( - config=uvicorn.Config( - app=app, - host=host, - port=port, - log_level="error", - lifespan="on", - ) - ) - server.run() - except Exception as e: - print(f"Server error: {e}") - sys.exit(1) - sys.exit(0) - - -def run_sse_server(host: str, port: int) -> None: - try: - app = fastmcp_server().http_app(transport="sse") - server = uvicorn.Server( - config=uvicorn.Config( - app=app, - host=host, - port=port, - log_level="error", - lifespan="on", - ) - ) - server.run() - except Exception as e: - print(f"Server error: {e}") - sys.exit(1) - sys.exit(0) +def run_server(host: str, port: int, **kwargs) -> None: + fastmcp_server().run(host=host, port=port, **kwargs) @pytest.fixture(autouse=True, scope="module") def shttp_server() -> Generator[str, None, None]: - with run_server_in_process(run_shttp_server) as url: + with run_server_in_process(run_server, transport="streamable-http") as url: yield f"{url}/mcp" @pytest.fixture(autouse=True, scope="module") def sse_server() -> Generator[str, None, None]: - with run_server_in_process(run_sse_server) as url: + with run_server_in_process(run_server, transport="sse") as url: yield f"{url}/sse"