gabotorresruiz commented on code in PR #44113: URL: https://github.com/apache/superset/pull/44113#discussion_r4008707457
########## tests/unit_tests/mcp_service/dashboard/tool/test_get_dashboard_data.py: ########## @@ -0,0 +1,450 @@ +# 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. + +"""Unit tests for the get_dashboard_data MCP tool.""" + +from typing import Any +from unittest.mock import AsyncMock, Mock, patch + +import pytest +from fastmcp import Client +from pydantic import ValidationError + +from superset.mcp_service.app import mcp +from superset.mcp_service.chart.schemas import ( + ChartData, + ChartError, + ChartQueryResult, + DataColumn, + PerformanceMetadata, +) +from superset.utils import json + +TOOL = "get_dashboard_data" +EXEC = "superset.mcp_service.dashboard.tool.get_dashboard_data.execute_chart_data" +DAO = "superset.daos.dashboard.DashboardDAO.get_by_id_or_slug" +PERM = "superset.mcp_service.auth.check_tool_permission" + + [email protected] +def mcp_server() -> object: + return mcp + + +def _slice(chart_id: int, name: str, viz: str = "table") -> Mock: + slc = Mock() + slc.id = chart_id + slc.slice_name = name + slc.viz_type = viz + return slc + + +def _dashboard( + dash_id: int, + title: str, + slices: list[Mock], + position_json: str | None = None, +) -> Mock: + dashboard = Mock() + dashboard.id = dash_id + dashboard.dashboard_title = title + dashboard.slices = slices + dashboard.position_json = position_json + return dashboard + + +_UNSET = object() + + +def _chart_data( + chart_id: int, + name: str, + viz: str = "table", + rows: list[dict[str, Any]] | None = None, + query_results: list[ChartQueryResult] | None = None, + total_rows: Any = _UNSET, +) -> ChartData: + rows = rows if rows is not None else [{"country": "US", "cnt": 10}] + resolved_total = len(rows) if total_rows is _UNSET else total_rows + return ChartData( + query_results=query_results, + chart_id=chart_id, + chart_name=name, + chart_type=viz, + columns=[ + DataColumn( + name="country", + display_name="Country", + data_type="VARCHAR", + sample_values=["US"], + null_count=0, + unique_count=1, + ) + ], + data=rows, + row_count=len(rows), + total_rows=resolved_total, + data_freshness=None, + summary=f"{name} summary", + insights=[f"{name} insight"], + data_quality={}, + recommended_visualizations=[], + performance=PerformanceMetadata(query_duration_ms=1, cache_status="miss"), + ) + + +async def _call(client: Client, request: dict[str, Any]) -> dict[str, Any]: + result = await client.call_tool(TOOL, {"request": request}) + return json.loads(result.content[0].text) + + +@patch(EXEC, new_callable=AsyncMock) +@patch(DAO) [email protected] +async def test_returns_compact_summary_per_chart(mock_dao, mock_exec, mcp_server): + mock_dao.return_value = _dashboard( + 6, "FCC Survey", [_slice(56, "Gender"), _slice(99, "Age")] + ) + mock_exec.side_effect = [_chart_data(56, "Gender"), _chart_data(99, "Age")] + + async with Client(mcp_server) as client: + data = await _call(client, {"identifier": 6}) + + assert data["dashboard_id"] == 6 + assert data["chart_count"] == 2 + assert data["charts_returned"] == 2 + assert data["charts_truncated"] is False + assert [c["chart_name"] for c in data["charts"]] == ["Gender", "Age"] + + first = data["charts"][0] + assert first["columns"] == ["country"] + assert first["sample_data"] == [{"country": "US", "cnt": 10}] + assert first["row_count"] == 1 + # Shallow summary/insights fields were dropped to keep the payload lean. + assert "summary" not in first + assert "insights" not in first + + +@patch(EXEC, new_callable=AsyncMock) +@patch(DAO) [email protected] +async def test_truncates_to_max_charts(mock_dao, mock_exec, mcp_server): + mock_dao.return_value = _dashboard( + 6, "D", [_slice(1, "a"), _slice(2, "b"), _slice(3, "c")] + ) + mock_exec.side_effect = [_chart_data(1, "a"), _chart_data(2, "b")] + + async with Client(mcp_server) as client: + data = await _call(client, {"identifier": 6, "max_charts": 2}) + + assert data["chart_count"] == 3 + assert data["charts_returned"] == 2 + assert data["charts_truncated"] is True + assert mock_exec.call_count == 2 + + +@patch(EXEC, new_callable=AsyncMock) +@patch(DAO) [email protected] +async def test_applies_filters_only_to_scoped_chart(mock_dao, mock_exec, mcp_server): + mock_dao.return_value = _dashboard( + 6, "D", [_slice(56, "Gender"), _slice(99, "Age")] + ) + mock_exec.side_effect = [_chart_data(56, "Gender"), _chart_data(99, "Age")] + gender_filter = {"filters": [{"col": "gender", "op": "IN", "val": ["Female"]}]} + + async with Client(mcp_server) as client: + data = await _call( + client, {"identifier": 6, "applied_filters": {"56": gender_filter}} + ) + + by_id = {c["chart_id"]: c for c in data["charts"]} + assert by_id[56]["filtered"] is True + assert by_id[99]["filtered"] is False + + # Only the scoped chart's query received the extra_form_data. + passed = { + call.args[0].identifier: call.args[0].extra_form_data + for call in mock_exec.call_args_list + } + assert passed[56] == gender_filter + assert passed[99] is None Review Comment: Pinned it: `assert {call.args[0].limit for call in mock_exec.call_args_list} == {100}`. Confirmed that deleting `limit=request.fetch_row_limit` now turns the test red. ########## tests/unit_tests/mcp_service/dashboard/tool/test_get_dashboard_data.py: ########## @@ -0,0 +1,450 @@ +# 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. + +"""Unit tests for the get_dashboard_data MCP tool.""" + +from typing import Any +from unittest.mock import AsyncMock, Mock, patch + +import pytest +from fastmcp import Client +from pydantic import ValidationError + +from superset.mcp_service.app import mcp +from superset.mcp_service.chart.schemas import ( + ChartData, + ChartError, + ChartQueryResult, + DataColumn, + PerformanceMetadata, +) +from superset.utils import json + +TOOL = "get_dashboard_data" +EXEC = "superset.mcp_service.dashboard.tool.get_dashboard_data.execute_chart_data" +DAO = "superset.daos.dashboard.DashboardDAO.get_by_id_or_slug" +PERM = "superset.mcp_service.auth.check_tool_permission" + + [email protected] +def mcp_server() -> object: + return mcp + + +def _slice(chart_id: int, name: str, viz: str = "table") -> Mock: + slc = Mock() + slc.id = chart_id + slc.slice_name = name + slc.viz_type = viz + return slc + + +def _dashboard( + dash_id: int, + title: str, + slices: list[Mock], + position_json: str | None = None, +) -> Mock: + dashboard = Mock() + dashboard.id = dash_id + dashboard.dashboard_title = title + dashboard.slices = slices + dashboard.position_json = position_json + return dashboard + + +_UNSET = object() + + +def _chart_data( + chart_id: int, + name: str, + viz: str = "table", + rows: list[dict[str, Any]] | None = None, + query_results: list[ChartQueryResult] | None = None, + total_rows: Any = _UNSET, +) -> ChartData: + rows = rows if rows is not None else [{"country": "US", "cnt": 10}] + resolved_total = len(rows) if total_rows is _UNSET else total_rows + return ChartData( + query_results=query_results, + chart_id=chart_id, + chart_name=name, + chart_type=viz, + columns=[ + DataColumn( + name="country", + display_name="Country", + data_type="VARCHAR", + sample_values=["US"], + null_count=0, + unique_count=1, + ) + ], + data=rows, + row_count=len(rows), + total_rows=resolved_total, + data_freshness=None, + summary=f"{name} summary", + insights=[f"{name} insight"], + data_quality={}, + recommended_visualizations=[], + performance=PerformanceMetadata(query_duration_ms=1, cache_status="miss"), + ) + + +async def _call(client: Client, request: dict[str, Any]) -> dict[str, Any]: + result = await client.call_tool(TOOL, {"request": request}) + return json.loads(result.content[0].text) + + +@patch(EXEC, new_callable=AsyncMock) +@patch(DAO) [email protected] +async def test_returns_compact_summary_per_chart(mock_dao, mock_exec, mcp_server): + mock_dao.return_value = _dashboard( + 6, "FCC Survey", [_slice(56, "Gender"), _slice(99, "Age")] + ) + mock_exec.side_effect = [_chart_data(56, "Gender"), _chart_data(99, "Age")] + + async with Client(mcp_server) as client: + data = await _call(client, {"identifier": 6}) + + assert data["dashboard_id"] == 6 + assert data["chart_count"] == 2 + assert data["charts_returned"] == 2 + assert data["charts_truncated"] is False + assert [c["chart_name"] for c in data["charts"]] == ["Gender", "Age"] + + first = data["charts"][0] + assert first["columns"] == ["country"] + assert first["sample_data"] == [{"country": "US", "cnt": 10}] + assert first["row_count"] == 1 + # Shallow summary/insights fields were dropped to keep the payload lean. + assert "summary" not in first + assert "insights" not in first + + +@patch(EXEC, new_callable=AsyncMock) +@patch(DAO) [email protected] +async def test_truncates_to_max_charts(mock_dao, mock_exec, mcp_server): + mock_dao.return_value = _dashboard( + 6, "D", [_slice(1, "a"), _slice(2, "b"), _slice(3, "c")] + ) + mock_exec.side_effect = [_chart_data(1, "a"), _chart_data(2, "b")] + + async with Client(mcp_server) as client: + data = await _call(client, {"identifier": 6, "max_charts": 2}) + + assert data["chart_count"] == 3 + assert data["charts_returned"] == 2 + assert data["charts_truncated"] is True + assert mock_exec.call_count == 2 + + +@patch(EXEC, new_callable=AsyncMock) +@patch(DAO) [email protected] +async def test_applies_filters_only_to_scoped_chart(mock_dao, mock_exec, mcp_server): + mock_dao.return_value = _dashboard( + 6, "D", [_slice(56, "Gender"), _slice(99, "Age")] + ) + mock_exec.side_effect = [_chart_data(56, "Gender"), _chart_data(99, "Age")] + gender_filter = {"filters": [{"col": "gender", "op": "IN", "val": ["Female"]}]} + + async with Client(mcp_server) as client: + data = await _call( + client, {"identifier": 6, "applied_filters": {"56": gender_filter}} + ) + + by_id = {c["chart_id"]: c for c in data["charts"]} + assert by_id[56]["filtered"] is True + assert by_id[99]["filtered"] is False + + # Only the scoped chart's query received the extra_form_data. + passed = { + call.args[0].identifier: call.args[0].extra_form_data + for call in mock_exec.call_args_list + } + assert passed[56] == gender_filter + assert passed[99] is None + + +@patch(EXEC, new_callable=AsyncMock) +@patch(DAO) [email protected] +async def test_chart_error_is_recorded_not_fatal(mock_dao, mock_exec, mcp_server): + mock_dao.return_value = _dashboard( + 6, "D", [_slice(56, "Gender"), _slice(99, "Age")] + ) + mock_exec.side_effect = [ + ChartError(error_type="DatasetNotAccessible", message="no access"), + _chart_data(99, "Age"), + ] + + async with Client(mcp_server) as client: + data = await _call(client, {"identifier": 6}) + + by_id = {c["chart_id"]: c for c in data["charts"]} + assert by_id[56]["error"] == "no access" + assert by_id[56]["sample_data"] == [] + assert by_id[99]["error"] is None + assert by_id[99]["sample_data"] == [{"country": "US", "cnt": 10}] + + +@patch(DAO) [email protected] +async def test_dashboard_not_found_returns_error(mock_dao, mcp_server): + mock_dao.side_effect = Exception("nope") + + async with Client(mcp_server) as client: + data = await _call(client, {"identifier": 999}) + + assert data["error_type"] == "DashboardNotFound" + assert "999" in data["error"] + + +@patch(EXEC, new_callable=AsyncMock) +@patch(DAO) [email protected] +async def test_selects_charts_in_layout_reading_order(mock_dao, mock_exec, mcp_server): + # Layout places chart 99 before chart 56, opposite the slice order. + position_json = json.dumps( + { + "ROOT_ID": {"id": "ROOT_ID", "type": "ROOT", "children": ["GRID_ID"]}, + "GRID_ID": { + "id": "GRID_ID", + "type": "GRID", + "children": ["CHART-b", "CHART-a"], + }, + "CHART-a": { + "id": "CHART-a", + "type": "CHART", + "meta": {"chartId": 56}, + "children": [], + }, + "CHART-b": { + "id": "CHART-b", + "type": "CHART", + "meta": {"chartId": 99}, + "children": [], + }, + } + ) + mock_dao.return_value = _dashboard( + 6, "D", [_slice(56, "Gender"), _slice(99, "Age")], position_json=position_json + ) + mock_exec.side_effect = [_chart_data(99, "Age")] + + async with Client(mcp_server) as client: + data = await _call(client, {"identifier": 6, "max_charts": 1}) + + # With max_charts=1, the layout-first chart (99) is the one selected. + assert data["charts_returned"] == 1 + assert data["charts"][0]["chart_id"] == 99 Review Comment: Good catch. Added `assert mock_exec.call_args.args[0].identifier == 99` so the assertion is on the slice actually queried, not the mocked result id. Verified it fails when `_order_by_layout` returns the slices unchanged or reversed. ########## tests/unit_tests/mcp_service/dashboard/tool/test_get_dashboard_data.py: ########## @@ -0,0 +1,450 @@ +# 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. + +"""Unit tests for the get_dashboard_data MCP tool.""" + +from typing import Any +from unittest.mock import AsyncMock, Mock, patch + +import pytest +from fastmcp import Client +from pydantic import ValidationError + +from superset.mcp_service.app import mcp +from superset.mcp_service.chart.schemas import ( + ChartData, + ChartError, + ChartQueryResult, + DataColumn, + PerformanceMetadata, +) +from superset.utils import json + +TOOL = "get_dashboard_data" +EXEC = "superset.mcp_service.dashboard.tool.get_dashboard_data.execute_chart_data" +DAO = "superset.daos.dashboard.DashboardDAO.get_by_id_or_slug" +PERM = "superset.mcp_service.auth.check_tool_permission" + + [email protected] +def mcp_server() -> object: + return mcp + + +def _slice(chart_id: int, name: str, viz: str = "table") -> Mock: + slc = Mock() + slc.id = chart_id + slc.slice_name = name + slc.viz_type = viz + return slc + + +def _dashboard( + dash_id: int, + title: str, + slices: list[Mock], + position_json: str | None = None, +) -> Mock: + dashboard = Mock() + dashboard.id = dash_id + dashboard.dashboard_title = title + dashboard.slices = slices + dashboard.position_json = position_json + return dashboard + + +_UNSET = object() + + +def _chart_data( + chart_id: int, + name: str, + viz: str = "table", + rows: list[dict[str, Any]] | None = None, + query_results: list[ChartQueryResult] | None = None, + total_rows: Any = _UNSET, +) -> ChartData: + rows = rows if rows is not None else [{"country": "US", "cnt": 10}] + resolved_total = len(rows) if total_rows is _UNSET else total_rows + return ChartData( + query_results=query_results, + chart_id=chart_id, + chart_name=name, + chart_type=viz, + columns=[ + DataColumn( + name="country", + display_name="Country", + data_type="VARCHAR", + sample_values=["US"], + null_count=0, + unique_count=1, + ) + ], + data=rows, + row_count=len(rows), + total_rows=resolved_total, + data_freshness=None, + summary=f"{name} summary", + insights=[f"{name} insight"], + data_quality={}, + recommended_visualizations=[], + performance=PerformanceMetadata(query_duration_ms=1, cache_status="miss"), + ) + + +async def _call(client: Client, request: dict[str, Any]) -> dict[str, Any]: + result = await client.call_tool(TOOL, {"request": request}) + return json.loads(result.content[0].text) + + +@patch(EXEC, new_callable=AsyncMock) +@patch(DAO) [email protected] +async def test_returns_compact_summary_per_chart(mock_dao, mock_exec, mcp_server): + mock_dao.return_value = _dashboard( + 6, "FCC Survey", [_slice(56, "Gender"), _slice(99, "Age")] + ) + mock_exec.side_effect = [_chart_data(56, "Gender"), _chart_data(99, "Age")] + + async with Client(mcp_server) as client: + data = await _call(client, {"identifier": 6}) + + assert data["dashboard_id"] == 6 + assert data["chart_count"] == 2 + assert data["charts_returned"] == 2 + assert data["charts_truncated"] is False + assert [c["chart_name"] for c in data["charts"]] == ["Gender", "Age"] + + first = data["charts"][0] + assert first["columns"] == ["country"] + assert first["sample_data"] == [{"country": "US", "cnt": 10}] + assert first["row_count"] == 1 + # Shallow summary/insights fields were dropped to keep the payload lean. + assert "summary" not in first + assert "insights" not in first + + +@patch(EXEC, new_callable=AsyncMock) +@patch(DAO) [email protected] +async def test_truncates_to_max_charts(mock_dao, mock_exec, mcp_server): + mock_dao.return_value = _dashboard( + 6, "D", [_slice(1, "a"), _slice(2, "b"), _slice(3, "c")] + ) + mock_exec.side_effect = [_chart_data(1, "a"), _chart_data(2, "b")] + + async with Client(mcp_server) as client: + data = await _call(client, {"identifier": 6, "max_charts": 2}) + + assert data["chart_count"] == 3 + assert data["charts_returned"] == 2 + assert data["charts_truncated"] is True + assert mock_exec.call_count == 2 + + +@patch(EXEC, new_callable=AsyncMock) +@patch(DAO) [email protected] +async def test_applies_filters_only_to_scoped_chart(mock_dao, mock_exec, mcp_server): + mock_dao.return_value = _dashboard( + 6, "D", [_slice(56, "Gender"), _slice(99, "Age")] + ) + mock_exec.side_effect = [_chart_data(56, "Gender"), _chart_data(99, "Age")] + gender_filter = {"filters": [{"col": "gender", "op": "IN", "val": ["Female"]}]} + + async with Client(mcp_server) as client: + data = await _call( + client, {"identifier": 6, "applied_filters": {"56": gender_filter}} + ) + + by_id = {c["chart_id"]: c for c in data["charts"]} + assert by_id[56]["filtered"] is True + assert by_id[99]["filtered"] is False + + # Only the scoped chart's query received the extra_form_data. + passed = { + call.args[0].identifier: call.args[0].extra_form_data + for call in mock_exec.call_args_list + } + assert passed[56] == gender_filter + assert passed[99] is None + + +@patch(EXEC, new_callable=AsyncMock) +@patch(DAO) [email protected] +async def test_chart_error_is_recorded_not_fatal(mock_dao, mock_exec, mcp_server): + mock_dao.return_value = _dashboard( + 6, "D", [_slice(56, "Gender"), _slice(99, "Age")] + ) + mock_exec.side_effect = [ + ChartError(error_type="DatasetNotAccessible", message="no access"), + _chart_data(99, "Age"), + ] + + async with Client(mcp_server) as client: + data = await _call(client, {"identifier": 6}) + + by_id = {c["chart_id"]: c for c in data["charts"]} + assert by_id[56]["error"] == "no access" + assert by_id[56]["sample_data"] == [] + assert by_id[99]["error"] is None + assert by_id[99]["sample_data"] == [{"country": "US", "cnt": 10}] Review Comment: Added `test_malformed_filter_is_recorded_not_fatal`: a non dict `applied_filters` value makes `GetChartDataRequest` raise inside the loop, which exercises the `except Exception` branch at `get_dashboard_data.py`. It asserts chart `56` gets an `error`, that `execute_chart_data` is never called for it, and that chart `99` still runs. Deleting the branch fails it. -- 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]
