From 02a905e8b95871367cdb2a6e903f2c5b3798867d Mon Sep 17 00:00:00 2001 From: Goro Date: Fri, 27 Jun 2025 20:04:25 +0200 Subject: [PATCH] Add component manager tests --- tests/contrib/test_component_manager.py | 537 ++++++++++++++++++++++++ 1 file changed, 537 insertions(+) create mode 100644 tests/contrib/test_component_manager.py diff --git a/tests/contrib/test_component_manager.py b/tests/contrib/test_component_manager.py new file mode 100644 index 000000000..1c36d42f9 --- /dev/null +++ b/tests/contrib/test_component_manager.py @@ -0,0 +1,537 @@ +import pytest +from starlette import status +from starlette.testclient import TestClient + +from fastmcp import FastMCP +from fastmcp.contrib.component_manager.component_manager import set_up_component_manager +from fastmcp.server.auth.providers.bearer import BearerAuthProvider, RSAKeyPair + + +class TestComponentManagementRoutes: + """Test the component management routes for tools, resources, and prompts.""" + + @pytest.fixture + def mounted_mcp(self): + """Create a FastMCP server with a mounted sub-server and a tool, resource, and prompt on the sub-server.""" + mounted_mcp = FastMCP("SubServer") + + @mounted_mcp.tool() + def mounted_tool() -> str: + """Test tool for tool management routes.""" + return "mounted_tool_result" + + @mounted_mcp.resource("data://mounted_resource") + def mounted_resource() -> str: + """Test resource for tool management routes.""" + return "mounted_resource_result" + + # Add a test resource + @mounted_mcp.resource("data://mounted_resource/{id}") + def test_template(id: str) -> dict: + """Test template for tool management routes.""" + return {"id": id, "value": "data"} + + @mounted_mcp.prompt() + def mounted_prompt() -> str: + """Test prompt for tool management routes.""" + return "mounted_prompt_result" + + return mounted_mcp + + @pytest.fixture + def mcp(self, mounted_mcp): + """Create a FastMCP server with test tools, resources, and prompts.""" + mcp = FastMCP("TestServer") + mcp.mount(mounted_mcp, prefix="sub") + set_up_component_manager(server=mcp) + # Add a test tool + @mcp.tool + def test_tool() -> str: + """Test tool for tool management routes.""" + return "test_tool_result" + + # Add a test resource + @mcp.resource("data://test_resource") + def test_resource() -> str: + """Test resource for tool management routes.""" + return "test_resource_result" + + # Add a test resource + @mcp.resource("data://test_resource/{id}") + def test_template(id: str) -> dict: + """Test template for tool management routes.""" + return {"id": id, "value": "data"} + + # Add a test prompt + @mcp.prompt + def test_prompt() -> str: + """Test prompt for tool management routes.""" + return "test_prompt_result" + + return mcp + + @pytest.fixture + def client(self, mcp): + """Create a test client for the FastMCP server.""" + return TestClient(mcp.http_app()) + + async def test_enable_tool_route(self, client, mcp): + """Test enabling a tool via the HTTP route.""" + # First disable the tool + tool = await mcp._tool_manager.get_tool("test_tool") + tool.enabled = False + + # Enable the tool via the HTTP route + response = client.post("/tools/test_tool/enable") + + assert response.status_code == status.HTTP_200_OK + assert response.json() == {"message": "Enabled tool: test_tool"} + + # Verify the tool is enabled + tool = await mcp._tool_manager.get_tool("test_tool") + assert tool.enabled is True + + async def test_disable_tool_route(self, client, mcp): + """Test disabling a tool via the HTTP route.""" + # First ensure the tool is enabled + tool = await mcp._tool_manager.get_tool("test_tool") + tool.enabled = True + + # Disable the tool via the HTTP route + response = client.post("/tools/test_tool/disable") + + assert response.status_code == status.HTTP_200_OK + assert response.json() == {"message": "Disabled tool: test_tool"} + + # Verify the tool is disabled + tool = await mcp._tool_manager.get_tool("test_tool") + assert tool.enabled is False + + async def test_enable_resource_route(self, client, mcp): + """Test enabling a resource via the HTTP route.""" + # First disable the resource + resource = await mcp._resource_manager.get_resource("data://test_resource") + resource.enabled = False + + # Enable the resource via the HTTP route + response = client.post("/resources/data://test_resource/enable") + + assert response.status_code == status.HTTP_200_OK + assert response.json() == {"message": "Enabled resource: data://test_resource"} + + # Verify the resource is enabled + resource = await mcp._resource_manager.get_resource("data://test_resource") + assert resource.enabled is True + + async def test_disable_resource_route(self, client, mcp): + """Test disabling a resource via the HTTP route.""" + # First ensure the resource is enabled + resource = await mcp._resource_manager.get_resource("data://test_resource") + resource.enabled = True + + # Disable the resource via the HTTP route + response = client.post("/resources/data://test_resource/disable") + + assert response.status_code == status.HTTP_200_OK + assert response.json() == {"message": "Disabled resource: data://test_resource"} + + # Verify the resource is disabled + resource = await mcp._resource_manager.get_resource("data://test_resource") + assert resource.enabled is False + + async def test_enable_template_route(self, client, mcp): + """Test enabling a resource on a mounted server via the parent server's HTTP route.""" + key = "data://test_resource/{id}" + resource = mcp._resource_manager._templates[key] + resource.enabled = False + response = client.post("/resources/data://test_resource/{id}/enable") + assert response.status_code == status.HTTP_200_OK + assert response.json() == { + "message": "Enabled resource: data://test_resource/{id}" + } + assert resource.enabled is True + + async def test_disable_template_route(self, client, mcp): + """Test disabling a resource on a mounted server via the parent server's HTTP route.""" + key = "data://test_resource/{id}" + resource = mcp._resource_manager._templates[key] + resource.enabled = True + response = client.post("/resources/data://test_resource/{id}/disable") + assert response.status_code == status.HTTP_200_OK + assert response.json() == { + "message": "Disabled resource: data://test_resource/{id}" + } + assert resource.enabled is False + + async def test_enable_prompt_route(self, client, mcp): + """Test enabling a prompt via the HTTP route.""" + # First disable the prompt + prompt = await mcp._prompt_manager.get_prompt("test_prompt") + prompt.enabled = False + + # Enable the prompt via the HTTP route + response = client.post("/prompts/test_prompt/enable") + + assert response.status_code == status.HTTP_200_OK + assert response.json() == {"message": "Enabled prompt: test_prompt"} + + # Verify the prompt is enabled + prompt = await mcp._prompt_manager.get_prompt("test_prompt") + assert prompt.enabled is True + + async def test_disable_prompt_route(self, client, mcp): + """Test disabling a prompt via the HTTP route.""" + # First ensure the prompt is enabled + prompt = await mcp._prompt_manager.get_prompt("test_prompt") + prompt.enabled = True + + # Disable the prompt via the HTTP route + response = client.post("/prompts/test_prompt/disable") + + assert response.status_code == status.HTTP_200_OK + assert response.json() == {"message": "Disabled prompt: test_prompt"} + + # Verify the prompt is disabled + prompt = await mcp._prompt_manager.get_prompt("test_prompt") + assert prompt.enabled is False + + async def test_enable_tool_route_on_mounted_server(self, client, mounted_mcp): + """Test enabling a tool on a mounted server via the parent server's HTTP route.""" + # Disable the tool on the sub-server + sub_tool = await mounted_mcp._tool_manager.get_tool("mounted_tool") + sub_tool.enabled = False + # Enable via parent + response = client.post("/tools/sub_mounted_tool/enable") + assert response.status_code == status.HTTP_200_OK + assert response.json() == {"message": "Enabled tool: sub_mounted_tool"} + # Confirm disabled on sub-server + assert sub_tool.enabled is True + + async def test_disable_tool_route_on_mounted_server(self, client, mounted_mcp): + """Test disabling a tool on a mounted server via the parent server's HTTP route.""" + # Enable the tool on the sub-server + sub_tool = await mounted_mcp._tool_manager.get_tool("mounted_tool") + sub_tool.enabled = True + # Disable via parent + response = client.post("/tools/sub_mounted_tool/disable") + assert response.status_code == status.HTTP_200_OK + assert response.json() == {"message": "Disabled tool: sub_mounted_tool"} + # Confirm disabled on sub-server + assert sub_tool.enabled is False + + async def test_enable_resource_route_on_mounted_server(self, client, mounted_mcp): + """Test enabling a resource on a mounted server via the parent server's HTTP route.""" + resource = await mounted_mcp._resource_manager.get_resource( + "data://mounted_resource" + ) + resource.enabled = False + response = client.post("/resources/data://sub/mounted_resource/enable") + assert response.status_code == status.HTTP_200_OK + assert response.json() == { + "message": "Enabled resource: data://sub/mounted_resource" + } + resource = await mounted_mcp._resource_manager.get_resource( + "data://mounted_resource" + ) + assert resource.enabled is True + + async def test_disable_resource_route_on_mounted_server(self, client, mounted_mcp): + """Test disabling a resource on a mounted server via the parent server's HTTP route.""" + resource = await mounted_mcp._resource_manager.get_resource( + "data://mounted_resource" + ) + resource.enabled = True + response = client.post("/resources/data://sub/mounted_resource/disable") + assert response.status_code == status.HTTP_200_OK + assert response.json() == { + "message": "Disabled resource: data://sub/mounted_resource" + } + resource = await mounted_mcp._resource_manager.get_resource( + "data://mounted_resource" + ) + assert resource.enabled is False + + async def test_enable_template_route_on_mounted_server(self, client, mounted_mcp): + """Test enabling a resource on a mounted server via the parent server's HTTP route.""" + key = "data://mounted_resource/{id}" + resource = mounted_mcp._resource_manager._templates[key] + resource.enabled = False + response = client.post("/resources/data://sub/mounted_resource/{id}/enable") + assert response.status_code == status.HTTP_200_OK + assert response.json() == { + "message": "Enabled resource: data://sub/mounted_resource/{id}" + } + assert resource.enabled is True + + async def test_disable_template_route_on_mounted_server(self, client, mounted_mcp): + """Test disabling a resource on a mounted server via the parent server's HTTP route.""" + key = "data://mounted_resource/{id}" + resource = mounted_mcp._resource_manager._templates[key] + resource.enabled = True + response = client.post("/resources/data://sub/mounted_resource/{id}/disable") + assert response.status_code == status.HTTP_200_OK + assert response.json() == { + "message": "Disabled resource: data://sub/mounted_resource/{id}" + } + assert resource.enabled is False + + async def test_enable_prompt_route_on_mounted_server(self, client, mounted_mcp): + """Test enabling a prompt on a mounted server via the parent server's HTTP route.""" + prompt = await mounted_mcp._prompt_manager.get_prompt("mounted_prompt") + prompt.enabled = False + response = client.post("/prompts/sub_mounted_prompt/enable") + assert response.status_code == status.HTTP_200_OK + assert response.json() == {"message": "Enabled prompt: sub_mounted_prompt"} + prompt = await mounted_mcp._prompt_manager.get_prompt("mounted_prompt") + assert prompt.enabled is True + + async def test_disable_prompt_route_on_mounted_server(self, client, mounted_mcp): + """Test disabling a prompt on a mounted server via the parent server's HTTP route.""" + prompt = await mounted_mcp._prompt_manager.get_prompt("mounted_prompt") + prompt.enabled = True + response = client.post("/prompts/sub_mounted_prompt/disable") + assert response.status_code == status.HTTP_200_OK + assert response.json() == {"message": "Disabled prompt: sub_mounted_prompt"} + prompt = await mounted_mcp._prompt_manager.get_prompt("mounted_prompt") + assert prompt.enabled is False + + def test_enable_nonexistent_tool(self, client): + """Test enabling a non-existent tool returns 404.""" + response = client.post("/tools/nonexistent_tool/enable") + assert response.status_code == status.HTTP_404_NOT_FOUND + assert response.text == "Unknown tool: nonexistent_tool" + + def test_disable_nonexistent_tool(self, client): + """Test disabling a non-existent tool returns 404.""" + response = client.post("/tools/nonexistent_tool/disable") + assert response.status_code == status.HTTP_404_NOT_FOUND + assert response.text == "Unknown tool: nonexistent_tool" + + def test_enable_nonexistent_resource(self, client): + """Test enabling a non-existent resource returns 404.""" + response = client.post("/resources/nonexistent://resource/enable") + assert response.status_code == status.HTTP_404_NOT_FOUND + assert response.text == "Unknown resource: nonexistent://resource" + + def test_disable_nonexistent_resource(self, client): + """Test disabling a non-existent resource returns 404.""" + response = client.post("/resources/nonexistent://resource/disable") + assert response.status_code == status.HTTP_404_NOT_FOUND + assert response.text == "Unknown resource: nonexistent://resource" + + def test_enable_nonexistent_prompt(self, client): + """Test enabling a non-existent prompt returns 404.""" + response = client.post("/prompts/nonexistent_prompt/enable") + assert response.status_code == status.HTTP_404_NOT_FOUND + assert response.text == "Unknown prompt: nonexistent_prompt" + + def test_disable_nonexistent_prompt(self, client): + """Test disabling a non-existent prompt returns 404.""" + response = client.post("/prompts/nonexistent_prompt/disable") + assert response.status_code == status.HTTP_404_NOT_FOUND + assert response.text == "Unknown prompt: nonexistent_prompt" + + +class TestAuthComponentManagementRoutes: + """Test the component management routes with authentication for tools, resources, and prompts.""" + + def setup_method(self): + """Set up test fixtures.""" + # Generate a key pair and create an auth provider + key_pair = RSAKeyPair.generate() + self.auth = BearerAuthProvider( + public_key=key_pair.public_key, + issuer="https://dev.example.com", + audience="my-dev-server", + ) + self.mcp = FastMCP("TestServerWithAuth", auth=self.auth) + set_up_component_manager(server=self.mcp, required_scopes=["tool:write", "tool:read"]) + self.token = key_pair.create_token( + subject="dev-user", + issuer="https://dev.example.com", + audience="my-dev-server", + scopes=["tool:write", "tool:read"], + ) + self.token_without_scopes = key_pair.create_token( + subject="dev-user", + issuer="https://dev.example.com", + audience="my-dev-server", + scopes=["tool:read"], + ) + + # Add test components + @self.mcp.tool + def test_tool() -> str: + """Test tool for auth testing.""" + return "test_tool_result" + + @self.mcp.resource("data://test_resource") + def test_resource() -> str: + """Test resource for auth testing.""" + return "test_resource_result" + + @self.mcp.prompt + def test_prompt() -> str: + """Test prompt for auth testing.""" + return "test_prompt_result" + + # Create test client + self.client = TestClient(self.mcp.http_app()) + + async def test_unauthorized_enable_tool(self): + """Test that unauthenticated requests to enable a tool are rejected.""" + tool = await self.mcp._tool_manager.get_tool("test_tool") + tool.enabled = False + + response = self.client.post("/tools/test_tool/enable") + assert response.status_code == 401 + assert tool.enabled is False + + async def test_authorized_enable_tool(self): + """Test that authenticated requests to enable a tool are allowed.""" + tool = await self.mcp._tool_manager.get_tool("test_tool") + tool.enabled = False + + response = self.client.post( + "/tools/test_tool/enable", headers={"Authorization": "Bearer " + self.token} + ) + assert response.status_code == 200 + assert response.json() == {"message": "Enabled tool: test_tool"} + assert tool.enabled is True + + async def test_unauthorized_disable_tool(self): + """Test that unauthenticated requests to disable a tool are rejected.""" + tool = await self.mcp._tool_manager.get_tool("test_tool") + tool.enabled = True + + response = self.client.post("/tools/test_tool/disable") + assert response.status_code == 401 + assert tool.enabled is True + + async def test_authorized_disable_tool(self): + """Test that authenticated requests to disable a tool are allowed.""" + tool = await self.mcp._tool_manager.get_tool("test_tool") + tool.enabled = True + + response = self.client.post( + "/tools/test_tool/disable", + headers={"Authorization": "Bearer " + self.token}, + ) + assert response.status_code == 200 + assert response.json() == {"message": "Disabled tool: test_tool"} + assert tool.enabled is False + + async def test_forbidden_enable_tool(self): + """Test that unauthenticated requests to enable a resource are rejected.""" + tool = await self.mcp._tool_manager.get_tool("test_tool") + tool.enabled = False + + response = self.client.post( + "/tools/test_tool/enable", + headers={"Authorization": "Bearer " + self.token_without_scopes}, + ) + assert response.status_code == 403 + assert tool.enabled is False + + async def test_authorized_enable_resource(self): + """Test that authenticated requests to enable a resource are allowed.""" + resource = await self.mcp._resource_manager.get_resource("data://test_resource") + resource.enabled = False + + response = self.client.post( + "/resources/data://test_resource/enable", + headers={"Authorization": "Bearer " + self.token}, + ) + assert response.status_code == 200 + assert response.json() == {"message": "Enabled resource: data://test_resource"} + assert resource.enabled is True + + async def test_unauthorized_disable_resource(self): + """Test that unauthenticated requests to disable a resource are rejected.""" + resource = await self.mcp._resource_manager.get_resource("data://test_resource") + resource.enabled = True + + response = self.client.post("/resources/data://test_resource/disable") + assert response.status_code == 401 + assert resource.enabled is True + + async def test_forbidden_enable_resource(self): + """Test that unauthenticated requests to enable a resource are rejected.""" + resource = await self.mcp._resource_manager.get_resource("data://test_resource") + resource.enabled = False + + response = self.client.post( + "/resources/data://test_resource/disable", + headers={"Authorization": "Bearer " + self.token_without_scopes}, + ) + assert response.status_code == 403 + assert resource.enabled is False + + async def test_authorized_disable_resource(self): + """Test that authenticated requests to disable a resource are allowed.""" + resource = await self.mcp._resource_manager.get_resource("data://test_resource") + resource.enabled = True + + response = self.client.post( + "/resources/data://test_resource/disable", + headers={"Authorization": "Bearer " + self.token}, + ) + assert response.status_code == 200 + assert response.json() == {"message": "Disabled resource: data://test_resource"} + assert resource.enabled is False + + async def test_unauthorized_enable_prompt(self): + """Test that unauthenticated requests to enable a prompt are rejected.""" + prompt = await self.mcp._prompt_manager.get_prompt("test_prompt") + prompt.enabled = False + + response = self.client.post("/prompts/test_prompt/enable") + assert response.status_code == 401 + assert prompt.enabled is False + + async def test_authorized_enable_prompt(self): + """Test that authenticated requests to enable a prompt are allowed.""" + prompt = await self.mcp._prompt_manager.get_prompt("test_prompt") + prompt.enabled = False + + response = self.client.post( + "/prompts/test_prompt/enable", + headers={"Authorization": "Bearer " + self.token}, + ) + assert response.status_code == 200 + assert response.json() == {"message": "Enabled prompt: test_prompt"} + assert prompt.enabled is True + + async def test_unauthorized_disable_prompt(self): + """Test that unauthenticated requests to disable a prompt are rejected.""" + prompt = await self.mcp._prompt_manager.get_prompt("test_prompt") + prompt.enabled = True + + response = self.client.post("/prompts/test_prompt/disable") + assert response.status_code == 401 + assert prompt.enabled is True + + async def test_forbidden_disable_prompt(self): + """Test that unauthenticated requests to enable a resource are rejected.""" + prompt = await self.mcp._prompt_manager.get_prompt("test_prompt") + prompt.enabled = True + + response = self.client.post( + "/prompts/test_prompt/disable", + headers={"Authorization": "Bearer " + self.token_without_scopes}, + ) + assert response.status_code == 403 + assert prompt.enabled is True + + async def test_authorized_disable_prompt(self): + """Test that authenticated requests to disable a prompt are allowed.""" + prompt = await self.mcp._prompt_manager.get_prompt("test_prompt") + prompt.enabled = True + + response = self.client.post( + "/prompts/test_prompt/disable", + headers={"Authorization": "Bearer " + self.token}, + ) + assert response.status_code == 200 + assert response.json() == {"message": "Disabled prompt: test_prompt"} + assert prompt.enabled is False \ No newline at end of file