From 6d5da3643aaa95b4ea8956aafd7b4fb8ba259a4f Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Sat, 3 May 2025 12:11:12 -0400
Subject: [PATCH] Update test_mount.py

---
 tests/server/test_mount.py | 109 +++++++++++++++++++++++++++++++++++++
 1 file changed, 109 insertions(+)

diff --git a/tests/server/test_mount.py b/tests/server/test_mount.py
index dedb9e802..1b495cc34 100644
--- a/tests/server/test_mount.py
+++ b/tests/server/test_mount.py
@@ -1,4 +1,5 @@
 import json
+from contextlib import asynccontextmanager
 
 import pytest
 from mcp.server.lowlevel.helper_types import ReadResourceContents
@@ -8,6 +9,7 @@ from fastmcp import FastMCP
 from fastmcp.client import Client
 from fastmcp.client.transports import FastMCPTransport
 from fastmcp.exceptions import NotFoundError
+from fastmcp.server.proxy import FastMCPProxy
 
 
 class TestBasicMount:
@@ -427,3 +429,110 @@ class TestProxyServer:
         result = await main_app._mcp_get_prompt("proxy_welcome", {"name": "World"})
         assert result.messages is not None
         # The message should contain our welcome text
+
+
+class TestAsProxyKwarg:
+    """Test the as_proxy kwarg."""
+
+    async def test_as_proxy_defaults_false(self):
+        mcp = FastMCP("Main")
+        sub = FastMCP("Sub")
+
+        mcp.mount("sub", sub)
+
+        assert mcp._mounted_servers["sub"].server is sub
+
+    async def test_as_proxy_false(self):
+        mcp = FastMCP("Main")
+        sub = FastMCP("Sub")
+
+        mcp.mount("sub", sub, as_proxy=False)
+
+        assert mcp._mounted_servers["sub"].server is sub
+
+    async def test_as_proxy_true(self):
+        mcp = FastMCP("Main")
+        sub = FastMCP("Sub")
+
+        mcp.mount("sub", sub, as_proxy=True)
+
+        assert mcp._mounted_servers["sub"].server is not sub
+        assert isinstance(mcp._mounted_servers["sub"].server, FastMCPProxy)
+
+    async def test_as_proxy_defaults_true_if_lifespan(self):
+        @asynccontextmanager
+        async def lifespan(mcp: FastMCP):
+            yield
+
+        mcp = FastMCP("Main")
+        sub = FastMCP("Sub", lifespan=lifespan)
+
+        mcp.mount("sub", sub)
+
+        assert mcp._mounted_servers["sub"].server is not sub
+        assert isinstance(mcp._mounted_servers["sub"].server, FastMCPProxy)
+
+    async def test_as_proxy_ignored_for_proxy_mounts_default(self):
+        mcp = FastMCP("Main")
+        sub = FastMCP("Sub")
+        sub_proxy = FastMCP.from_client(Client(transport=FastMCPTransport(sub)))
+
+        mcp.mount("sub", sub_proxy)
+
+        assert mcp._mounted_servers["sub"].server is sub_proxy
+
+    async def test_as_proxy_ignored_for_proxy_mounts_false(self):
+        mcp = FastMCP("Main")
+        sub = FastMCP("Sub")
+        sub_proxy = FastMCP.from_client(Client(transport=FastMCPTransport(sub)))
+
+        mcp.mount("sub", sub_proxy, as_proxy=False)
+
+        assert mcp._mounted_servers["sub"].server is sub_proxy
+
+    async def test_as_proxy_ignored_for_proxy_mounts_true(self):
+        mcp = FastMCP("Main")
+        sub = FastMCP("Sub")
+        sub_proxy = FastMCP.from_client(Client(transport=FastMCPTransport(sub)))
+
+        mcp.mount("sub", sub_proxy, as_proxy=True)
+
+        assert mcp._mounted_servers["sub"].server is sub_proxy
+
+    async def test_as_proxy_mounts_still_have_live_link(self):
+        mcp = FastMCP("Main")
+        sub = FastMCP("Sub")
+
+        mcp.mount("sub", sub, as_proxy=True)
+
+        assert len(await mcp.get_tools()) == 0
+
+        @sub.tool()
+        def hello():
+            return "hi"
+
+        assert len(await mcp.get_tools()) == 1
+
+    async def test_sub_lifespan_is_executed(self):
+        lifespan_check = []
+
+        @asynccontextmanager
+        async def lifespan(mcp: FastMCP):
+            lifespan_check.append("start")
+            yield
+
+        mcp = FastMCP("Main")
+        sub = FastMCP("Sub", lifespan=lifespan)
+
+        @sub.tool()
+        def hello():
+            return "hi"
+
+        mcp.mount("sub", sub, as_proxy=True)
+
+        assert lifespan_check == []
+
+        async with Client(mcp) as client:
+            await client.call_tool("sub_hello", {})
+
+        assert lifespan_check == ["start"]