Merge pull request #649 from jlowin/simplify-process

Simplify code for running servers in processes during tests

src/fastmcp/utilities/tests.py

@@ -72,14 +72,21 @@ 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: 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 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 function that runs a FastMCP server. FastMCP servers are
+            not pickleable, so we need a function that creates and runs one.
+        *args: Arguments to pass to the server function.
+        provide_host_and_port: Whether to provide the host and port to the server function as kwargs.
+        **kwargs: Keyword arguments to pass to the server function.
 
     Returns:
         The server URL.
@@ -87,8 +94,11 @@ def run_server_in_process(
     host = "127.0.0.1"
     port = find_available_port()
 
+    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()
 

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"
 
 

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
 
 
-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)
+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:
     @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):

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

tests/client/test_streamable_http.py

@@ -63,28 +63,33 @@ 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",
-            )
+def run_server(host: str, port: int, **kwargs) -> None:
+    fastmcp_server().run(host=host, port=port, **kwargs)
+
+
+def run_nested_server(host: str, port: int) -> None:
+    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,
+    )
+    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()
 
 
 @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"
 
 
@@ -110,31 +115,6 @@ async def test_http_headers(streamable_http_server: str):
         assert json_result["x-demo-header"] == "ABC"
 
 
-def run_nested_server(host: str, port: int) -> None:
-    try:
-        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,
-        )
-        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)
-
-
 async def test_nested_streamable_http_server_resolves_correctly():
     # tests patch for
     # https://github.com/modelcontextprotocol/python-sdk/pull/659

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"