"""Tests for search transforms.""" from __future__ import annotations from collections.abc import Sequence from typing import Any from unittest.mock import MagicMock import mcp_types import pytest from mcp_types import TextContent from fastmcp import Client, FastMCP from fastmcp.server.context import Context from fastmcp.server.middleware.middleware import CallNext, Middleware, MiddlewareContext from fastmcp.server.transforms import Visibility from fastmcp.server.transforms.search.bm25 import ( BM25SearchTransform, _BM25Index, _catalog_hash, ) from fastmcp.server.transforms.search.regex import RegexSearchTransform from fastmcp.tools.base import Tool, ToolResult # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _parse_tool_result(result: ToolResult) -> list[dict[str, Any]]: """Extract tool list from a ToolResult's structured content.""" assert result.structured_content is not None return result.structured_content["result"] def _make_server_with_tools() -> FastMCP: mcp = FastMCP("test") @mcp.tool def add(a: int, b: int) -> int: """Add two numbers together.""" return a + b @mcp.tool def multiply(x: float, y: float) -> float: """Multiply two numbers.""" return x * y @mcp.tool def search_database(query: str, limit: int = 10) -> str: """Search the database for records matching the query.""" return f"results for {query}" @mcp.tool def delete_record(record_id: str) -> str: """Delete a record from the database by its ID.""" return f"deleted {record_id}" @mcp.tool def send_email(to: str, subject: str, body: str) -> str: """Send an email to the given recipient.""" return "sent" return mcp # --------------------------------------------------------------------------- # Shared behavior tests (parameterized across both transforms) # --------------------------------------------------------------------------- class TestBaseTransformBehavior: """Tests for behavior shared by all search transforms.""" async def test_list_tools_hides_tools_regex(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) tools = await mcp.list_tools() names = {t.name for t in tools} assert names == {"search_tools", "call_tool"} async def test_list_tools_hides_tools_bm25(self): mcp = _make_server_with_tools() mcp.add_transform(BM25SearchTransform()) tools = await mcp.list_tools() names = {t.name for t in tools} assert names == {"search_tools", "call_tool"} async def test_always_visible_pins_tools(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform(always_visible=["add"])) tools = await mcp.list_tools() names = {t.name for t in tools} assert "add" in names assert "search_tools" in names assert "call_tool" in names assert "multiply" not in names async def test_get_tool_returns_synthetic(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) search = await mcp.get_tool("search_tools") assert search is not None assert search.name == "search_tools" call = await mcp.get_tool("call_tool") assert call is not None assert call.name == "call_tool" async def test_get_tool_passes_through_hidden(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) tool = await mcp.get_tool("add") assert tool is not None assert tool.name == "add" async def test_call_tool_proxy_executes(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) # Need to call list_tools to populate catalog await mcp.list_tools() result = await mcp.call_tool( "call_tool", {"name": "add", "arguments": {"a": 2, "b": 3}} ) assert any("5" in c.text for c in result.content if isinstance(c, TextContent)) async def test_custom_tool_names(self): mcp = _make_server_with_tools() mcp.add_transform( RegexSearchTransform( search_tool_name="find_tools", call_tool_name="run_tool", ) ) tools = await mcp.list_tools() names = {t.name for t in tools} assert names == {"find_tools", "run_tool"} assert await mcp.get_tool("find_tools") is not None assert await mcp.get_tool("run_tool") is not None async def test_search_respects_visibility_filtering(self): """Tools disabled via Visibility transform should not appear in search.""" mcp = _make_server_with_tools() mcp.add_transform(Visibility(False, names={"delete_record"})) mcp.add_transform(RegexSearchTransform()) tools = await mcp.list_tools() names = {t.name for t in tools} assert "delete_record" not in names result = await mcp.call_tool("search_tools", {"pattern": "delete"}) found = _parse_tool_result(result) assert not any(t["name"] == "delete_record" for t in found) async def test_search_respects_auth_middleware(self): """Tools filtered by auth middleware should not appear in search.""" class BlockAdminTools(Middleware): async def on_list_tools( self, context: MiddlewareContext[mcp_types.ListToolsRequest], call_next: CallNext[mcp_types.ListToolsRequest, Sequence[Tool]], ) -> Sequence[Tool]: tools = await call_next(context) return [t for t in tools if t.name != "delete_record"] mcp = _make_server_with_tools() mcp.add_middleware(BlockAdminTools()) mcp.add_transform(RegexSearchTransform()) async with Client(mcp) as client: tools = await client.list_tools() names = {t.name for t in tools} assert "delete_record" not in names assert "search_tools" in names result = await client.call_tool("search_tools", {"pattern": "delete"}) found = _parse_tool_result(result) assert not any(t["name"] == "delete_record" for t in found) async def test_search_respects_session_visibility(self): """Tools disabled via session visibility should not appear in search.""" mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) @mcp.tool async def disable_delete(ctx: Context) -> str: """Helper tool to disable delete_record for this session.""" await ctx.disable_components(names={"delete_record"}) return "disabled" async with Client(mcp) as client: # Before disabling, search should find delete_record result = await client.call_tool("search_tools", {"pattern": "delete"}) found = _parse_tool_result(result) assert any(t["name"] == "delete_record" for t in found) # Disable via session visibility await client.call_tool("disable_delete", {}) # After disabling, search should NOT find it result = await client.call_tool("search_tools", {"pattern": "delete"}) found = _parse_tool_result(result) assert not any(t["name"] == "delete_record" for t in found) # --------------------------------------------------------------------------- # Regex-specific tests # --------------------------------------------------------------------------- class TestRegexSearch: async def test_search_by_name(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) await mcp.list_tools() result = await mcp.call_tool("search_tools", {"pattern": "add"}) tools = _parse_tool_result(result) assert any(t["name"] == "add" for t in tools) async def test_search_by_description(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) await mcp.list_tools() result = await mcp.call_tool("search_tools", {"pattern": "email"}) tools = _parse_tool_result(result) assert any(t["name"] == "send_email" for t in tools) async def test_search_by_param_name(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) await mcp.list_tools() result = await mcp.call_tool("search_tools", {"pattern": "record_id"}) tools = _parse_tool_result(result) assert any(t["name"] == "delete_record" for t in tools) async def test_search_by_param_description(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) await mcp.list_tools() result = await mcp.call_tool("search_tools", {"pattern": "recipient"}) tools = _parse_tool_result(result) assert any(t["name"] == "send_email" for t in tools) async def test_search_or_pattern(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform(max_results=10)) await mcp.list_tools() result = await mcp.call_tool("search_tools", {"pattern": "add|multiply"}) tools = _parse_tool_result(result) names = {t["name"] for t in tools} assert "add" in names assert "multiply" in names async def test_search_case_insensitive(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) await mcp.list_tools() result = await mcp.call_tool("search_tools", {"pattern": "ADD"}) tools = _parse_tool_result(result) assert any(t["name"] == "add" for t in tools) async def test_search_invalid_pattern(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) await mcp.list_tools() result = await mcp.call_tool("search_tools", {"pattern": "[invalid"}) tools = _parse_tool_result(result) assert tools == [] async def test_search_max_results(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform(max_results=2)) await mcp.list_tools() # Match everything result = await mcp.call_tool("search_tools", {"pattern": ".*"}) tools = _parse_tool_result(result) assert len(tools) == 2 async def test_search_no_matches(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) await mcp.list_tools() result = await mcp.call_tool("search_tools", {"pattern": "zzz_nonexistent"}) tools = _parse_tool_result(result) assert tools == [] async def test_search_returns_full_schema(self): mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) await mcp.list_tools() result = await mcp.call_tool("search_tools", {"pattern": "add"}) tools = _parse_tool_result(result) add_tool = next(t for t in tools if t["name"] == "add") assert "inputSchema" in add_tool assert "properties" in add_tool["inputSchema"] # --------------------------------------------------------------------------- # BM25-specific tests # --------------------------------------------------------------------------- class TestBM25Search: async def test_search_relevance(self): mcp = _make_server_with_tools() mcp.add_transform(BM25SearchTransform()) await mcp.list_tools() result = await mcp.call_tool("search_tools", {"query": "database"}) tools = _parse_tool_result(result) # Database tools should rank highest assert len(tools) > 0 names = {t["name"] for t in tools} assert "search_database" in names async def test_search_database_tools(self): mcp = _make_server_with_tools() mcp.add_transform(BM25SearchTransform()) await mcp.list_tools() result = await mcp.call_tool( "search_tools", {"query": "delete records from database"} ) tools = _parse_tool_result(result) assert len(tools) > 0 # delete_record should be highly relevant assert tools[0]["name"] == "delete_record" async def test_search_max_results(self): mcp = _make_server_with_tools() mcp.add_transform(BM25SearchTransform(max_results=2)) await mcp.list_tools() result = await mcp.call_tool("search_tools", {"query": "number"}) tools = _parse_tool_result(result) assert len(tools) <= 2 async def test_search_no_matches(self): mcp = _make_server_with_tools() mcp.add_transform(BM25SearchTransform()) await mcp.list_tools() result = await mcp.call_tool("search_tools", {"query": "zzz_nonexistent_xyz"}) tools = _parse_tool_result(result) assert tools == [] async def test_index_rebuilds_on_catalog_change(self): """When a new tool is added, the next list_tools + search sees it.""" mcp = _make_server_with_tools() mcp.add_transform(BM25SearchTransform()) await mcp.list_tools() # Search before adding tool result = await mcp.call_tool("search_tools", {"query": "weather forecast"}) tools = _parse_tool_result(result) assert not any(t["name"] == "get_weather" for t in tools) # Add a new tool @mcp.tool def get_weather(city: str) -> str: """Get the weather forecast for a city.""" return f"sunny in {city}" # Must call list_tools to refresh the catalog cache await mcp.list_tools() result = await mcp.call_tool("search_tools", {"query": "weather forecast"}) tools = _parse_tool_result(result) assert any(t["name"] == "get_weather" for t in tools) async def test_search_empty_query(self): mcp = _make_server_with_tools() mcp.add_transform(BM25SearchTransform()) await mcp.list_tools() result = await mcp.call_tool("search_tools", {"query": ""}) tools = _parse_tool_result(result) assert tools == [] async def test_search_returns_full_schema(self): mcp = _make_server_with_tools() mcp.add_transform(BM25SearchTransform()) await mcp.list_tools() result = await mcp.call_tool("search_tools", {"query": "add numbers"}) tools = _parse_tool_result(result) add_tool = next(t for t in tools if t["name"] == "add") assert "inputSchema" in add_tool assert "properties" in add_tool["inputSchema"] # --------------------------------------------------------------------------- # BM25 index unit tests # --------------------------------------------------------------------------- class TestBM25Index: def test_basic_ranking(self): index = _BM25Index() index.build( [ "search database query records", "add two numbers together", "send email recipient subject", ] ) results = index.query("database records", 3) assert results[0] == 0 # Database doc should rank first def test_empty_corpus(self): index = _BM25Index() index.build([]) assert index.query("anything", 5) == [] def test_no_matching_tokens(self): index = _BM25Index() index.build(["alpha beta gamma"]) assert index.query("zzz", 5) == [] # --------------------------------------------------------------------------- # call_tool self-reference guard # --------------------------------------------------------------------------- class TestCallToolGuard: async def test_call_tool_proxy_rejects_itself(self): """Calling call_tool(name='call_tool') must not recurse infinitely.""" mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) async with Client(mcp) as client: with pytest.raises(Exception): await client.call_tool( "call_tool", {"name": "call_tool", "arguments": {}} ) async def test_call_tool_proxy_rejects_search_tool(self): """Calling call_tool(name='search_tools') must be rejected.""" mcp = _make_server_with_tools() mcp.add_transform(RegexSearchTransform()) async with Client(mcp) as client: with pytest.raises(Exception): await client.call_tool( "call_tool", {"name": "search_tools", "arguments": {"pattern": "add"}}, ) async def test_call_tool_proxy_rejects_custom_names(self): """Guard works when synthetic tools have custom names.""" mcp = _make_server_with_tools() mcp.add_transform( RegexSearchTransform( search_tool_name="find_tools", call_tool_name="run_tool" ) ) async with Client(mcp) as client: with pytest.raises(Exception): await client.call_tool( "run_tool", {"name": "run_tool", "arguments": {}} ) with pytest.raises(Exception): await client.call_tool( "run_tool", {"name": "find_tools", "arguments": {"pattern": "add"}} ) # --------------------------------------------------------------------------- # catalog hash staleness # --------------------------------------------------------------------------- class TestCatalogHash: def test_hash_differs_for_same_name_different_description(self): """Hash must change when a tool's description changes, not just its name.""" tool_a = MagicMock() tool_a.name = "search" tool_a.description = "find records in the database" tool_a.parameters = {} tool_b = MagicMock() tool_b.name = "search" tool_b.description = "send an email to a recipient" tool_b.parameters = {} assert _catalog_hash([tool_a]) != _catalog_hash([tool_b])