aminghadersohi commented on code in PR #45818: URL: https://github.com/apache/superset/pull/45818#discussion_r4404683888
########## tests/unit_tests/mcp_service/test_tool_search_structured_output.py: ########## @@ -0,0 +1,188 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""``search_tools`` honors its advertised output schema for restricted callers. + +``search_tools`` is typed ``str | list[dict]``, so FastMCP advertises an +``outputSchema`` for it whenever structured output is enabled. MCP clients +reject a successful call that omits ``structuredContent`` or returns content +that does not satisfy that schema ("Tool search_tools has an output schema but +did not return structured content"). A caller with few permissions sees a +mostly or fully filtered catalog, so these tests drive ``search_tools`` over +the real streamable-HTTP stack as such callers and validate the wire result +against the schema advertised in ``tools/list``. +""" + +from collections.abc import Callable, Iterator +from contextlib import ExitStack +from unittest.mock import Mock, patch + +import jsonschema +import pytest +from fastmcp import FastMCP +from flask import Flask + +from superset.mcp_service.app import mcp +from superset.mcp_service.mcp_config import MCP_TOOL_SEARCH_CONFIG +from superset.mcp_service.server import ( + _apply_tool_search_transform, + build_middleware_list, +) +from tests.unit_tests.mcp_service.test_mcp_e2e_smoke import _asgi_client + +SEARCH_TOOL = "search_tools" + +# Permissions a viewer limited to dashboards holds: dashboard and chart reads, +# plus the dataset read that dashboard rendering needs. No write access, no +# SQL Lab, and no dataset-metadata access. +DASHBOARDS_ONLY_PERMISSIONS = frozenset( + { + ("can_read", "Dashboard"), + ("can_read", "Chart"), + ("can_read", "Dataset"), + } +) +NO_PERMISSIONS: frozenset[tuple[str, str]] = frozenset() + +# Tools a dashboards-only viewer must never be offered. +WRITE_AND_SQL_TOOLS = { + "generate_chart", + "delete_chart", + "update_dashboard", + "execute_sql", +} + + [email protected] +def rbac_enabled(app: Flask) -> Iterator[None]: + """Enforce RBAC; the package-wide autouse fixture disables it.""" + with patch.dict(app.config, {"MCP_RBAC_ENABLED": True}): + yield + + +async def _build_server(strategy: str, *, structured_output_enabled: bool) -> FastMCP: + """Serve the real tool catalog behind the production middleware and search. + + The shared ``mcp`` singleton is left untouched: its tools are re-registered + on a throwaway server so the search transform and middleware cannot leak + into other tests. + """ + server = FastMCP("tool-search-structured-output") + for tool in await mcp.list_tools(): + server.add_tool(tool) + for middleware in build_middleware_list( + structured_output_enabled=structured_output_enabled + ): + server.add_middleware(middleware) + _apply_tool_search_transform( + server, {**MCP_TOOL_SEARCH_CONFIG, "strategy": strategy} + ) + return server + + [email protected] +def as_caller() -> Iterator[Callable[[frozenset[tuple[str, str]]], None]]: + """Return a function that authenticates the request as a permissioned caller.""" + with ExitStack() as stack: + + def _as_caller(granted: frozenset[tuple[str, str]]) -> None: + user = Mock() + user.id = 1 + user.username = "viewer" + stack.enter_context( + patch( + "superset.mcp_service.auth.get_user_from_request", + return_value=user, + ) + ) + # Tool visibility reads the manager through ``auth``, and the + # data-model privacy gate reads it from ``superset``. + for target in ( + "superset.mcp_service.auth.security_manager", + "superset.security_manager", + ): + manager = stack.enter_context(patch(target)) + manager.can_access = ( + lambda permission, view: ( + permission, + view, + ) + in granted + ) + manager.is_guest_user = lambda: False + + yield _as_caller + + [email protected] [email protected]("strategy", ["bm25", "regex"]) [email protected]("query", [None, "", "bar chart", "no such capability zzz"]) [email protected]( + "granted", [DASHBOARDS_ONLY_PERMISSIONS, NO_PERMISSIONS], ids=["dashboards", "none"] +) +async def test_search_tools_returns_schema_valid_structured_content( + rbac_enabled: None, + as_caller: Callable[[frozenset[tuple[str, str]]], None], + strategy: str, + query: str | None, + granted: frozenset[tuple[str, str]], +) -> None: + """A restricted caller gets structured content that satisfies the schema.""" + as_caller(granted) + server = await _build_server(strategy, structured_output_enabled=True) + arguments = {} if query is None else {"query": query} + + async with _asgi_client(server) as client: + listed = {tool.name: tool for tool in await client.list_tools()} + result = await client.call_tool(SEARCH_TOOL, arguments) + + output_schema = listed[SEARCH_TOOL].outputSchema + assert output_schema is not None, "structured output is enabled" + assert result.structured_content is not None, ( + "search_tools advertised an outputSchema but returned no structuredContent" + ) + jsonschema.validate(result.structured_content, output_schema) + returned = result.structured_content["result"] + assert isinstance(returned, list) Review Comment: Fixed in f85580a2c4361f0a6b604e14c6e9c72d2bd9c13b. I replaced the purported no-match phrase with `zzzxqvnosuchcapability` (the original phrase matches BM25 tokens) and asserted `returned == []` in `test_search_tools_returns_schema_valid_structured_content` for both search strategies and both restricted callers. Regression evidence: making search return the visible catalog instead of searching let all four no-match cases pass without the assertion; adding the assertion made all four fail. With production search restored, all 18 tests pass on FastMCP 4.1.2; the same mutation fails all four no-match cases there too. Siblings checked: - tests/unit_tests/mcp_service/test_tool_search_structured_output.py — fixed in f85580a2c4361f0a6b604e14c6e9c72d2bd9c13b; the structured-output-disabled sibling uses a positive query, so no no-match assertion applies. - tests/unit_tests/mcp_service/test_tool_search_transform.py — not affected: the three `no_match` parametrizations compare the complete result sequence with the upstream strategy, rather than only checking its type. - tests/unit_tests/mcp_service/test_native_tool_surface.py — not affected: search discovery uses an omitted query and compares visibility across native and compatibility surfaces. - tests/unit_tests/mcp_service/catalog/tool/test_get_catalog.py — not affected: tests a positive exact-name query and requires the matching entry. - tests/unit_tests/mcp_service/test_middleware_logging.py — not affected: tests name resolution, not search results. - tests/unit_tests/mcp_service/test_error_classification.py — not affected: tests proxy errors, not search results. - tests/unit_tests/mcp_service/semantic_layer/tool/test_list_metrics.py — not affected: no-match and nonexistent-filter cases already assert `metrics == []` and zero count. - tests/unit_tests/mcp_service/theme/tool/test_list_themes.py — not affected: the empty-result test already asserts `themes == []` and zero count. - tests/unit_tests/mcp_service/dataset/tool/test_dataset_tools.py — not affected: separate DatasetDAO substring-search coverage asserts zero matching count, not tool-catalog search. Validation: `pytest -q -n 4 tests/unit_tests/mcp_service/ -m "not manual and not llm_compliance"`: 13,317 passed, 12 skipped; changed-file `pre-commit run` passed, including mypy. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
