aminghadersohi commented on code in PR #40343: URL: https://github.com/apache/superset/pull/40343#discussion_r3312986933
########## tests/unit_tests/mcp_service/css_template/tool/test_css_template_tools.py: ########## @@ -0,0 +1,237 @@ +# 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. + +import logging +from unittest.mock import MagicMock, patch + +import pytest +from fastmcp import Client +from fastmcp.exceptions import ToolError +from pydantic import ValidationError + +from superset.mcp_service.app import mcp +from superset.mcp_service.css_template.schemas import ( + CssTemplateFilter, + ListCssTemplatesRequest, +) +from superset.utils import json + +logging.basicConfig(level=logging.DEBUG) +logger = logging.getLogger(__name__) + + +class TestCssTemplateFilterSchema: + """Tests for CssTemplateFilter schema — filterable columns.""" + + def test_invalid_filter_column_rejected(self): + """Columns not in the Literal set must be rejected.""" + with pytest.raises(ValidationError): + CssTemplateFilter(col="not_a_real_column", opr="eq", value="test") + + def test_valid_filter_column_accepted(self): + """template_name is a valid filter column.""" + f = CssTemplateFilter(col="template_name", opr="eq", value="my_template") + assert f.col == "template_name" + + def test_css_column_not_filterable(self): + """css is not a public filter column (large field).""" + with pytest.raises(ValidationError): + CssTemplateFilter(col="css", opr="eq", value="body {}") + + +def create_mock_css_template( + template_id: int = 1, + template_name: str = "my_template", + css: str = "body { color: red; }", +) -> MagicMock: + """Factory function to create mock CSS template objects.""" + template = MagicMock() + template.id = template_id + template.template_name = template_name + template.css = css + template.uuid = f"test-css-template-uuid-{template_id}" + template.changed_on = None + template.created_on = None + return template + + [email protected] +def mcp_server(): + return mcp + + [email protected](autouse=True) +def mock_auth(): + """Mock authentication for all tests.""" + from unittest.mock import Mock, patch + + with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user: + mock_user = Mock() + mock_user.id = 1 + mock_user.username = "admin" + mock_get_user.return_value = mock_user + yield mock_get_user + + +@patch("superset.daos.css.CssTemplateDAO.list") [email protected] +async def test_list_css_templates_basic(mock_list, mcp_server): + """Test basic CSS template listing functionality.""" + template = create_mock_css_template() + mock_list.return_value = ([template], 1) + + async with Client(mcp_server) as client: + request = ListCssTemplatesRequest(page=1, page_size=10) + result = await client.call_tool( + "list_css_templates", {"request": request.model_dump()} + ) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["css_templates"] is not None + assert len(data["css_templates"]) == 1 + assert data["css_templates"][0]["id"] == 1 + assert data["css_templates"][0]["template_name"] == "my_template" + + +@patch("superset.daos.css.CssTemplateDAO.list") [email protected] +async def test_list_css_templates_css_not_in_default_columns(mock_list, mcp_server): + """Test that css field is excluded from default columns (it's large).""" + template = create_mock_css_template() + mock_list.return_value = ([template], 1) + + async with Client(mcp_server) as client: + result = await client.call_tool("list_css_templates", {}) + data = json.loads(result.content[0].text) + + assert data["css_templates"] is not None + assert len(data["css_templates"]) == 1 + # css not in default columns + assert "css" not in data["css_templates"][0] + assert data["columns_requested"] == ["id", "uuid", "template_name"] + assert data["css_templates"][0]["uuid"] == "test-css-template-uuid-1" + + +@patch("superset.daos.css.CssTemplateDAO.list") [email protected] +async def test_list_css_templates_with_search(mock_list, mcp_server): + """Test CSS template listing with search functionality.""" + template = create_mock_css_template(template_name="dark_theme_css") + mock_list.return_value = ([template], 1) + + async with Client(mcp_server) as client: + request = ListCssTemplatesRequest(page=1, page_size=10, search="dark") + result = await client.call_tool( + "list_css_templates", {"request": request.model_dump()} + ) + data = json.loads(result.content[0].text) + assert len(data["css_templates"]) == 1 + assert data["css_templates"][0]["template_name"] == "dark_theme_css" + + +@patch("superset.daos.css.CssTemplateDAO.list") [email protected] +async def test_list_css_templates_with_select_columns_css(mock_list, mcp_server): + """Test that css can be requested via select_columns.""" + template = create_mock_css_template(css="body { margin: 0; }") + mock_list.return_value = ([template], 1) + + async with Client(mcp_server) as client: + request = ListCssTemplatesRequest( + page=1, page_size=10, select_columns=["id", "template_name", "css"] + ) + result = await client.call_tool( + "list_css_templates", {"request": request.model_dump()} + ) + data = json.loads(result.content[0].text) + assert data["css_templates"][0]["css"] == "body { margin: 0; }" + + +@patch("superset.daos.css.CssTemplateDAO.list") [email protected] +async def test_list_css_templates_with_filters(mock_list, mcp_server): + """Test CSS template listing with filters.""" + template = create_mock_css_template(template_name="bootstrap_css") + mock_list.return_value = ([template], 1) + + async with Client(mcp_server) as client: + request = ListCssTemplatesRequest( + page=1, + page_size=10, + filters=[{"col": "template_name", "opr": "eq", "value": "bootstrap_css"}], + ) + result = await client.call_tool( + "list_css_templates", {"request": request.model_dump()} + ) + data = json.loads(result.content[0].text) + assert len(data["css_templates"]) == 1 + + +@patch("superset.daos.css.CssTemplateDAO.list") [email protected] +async def test_list_css_templates_api_error(mock_list, mcp_server): + """Test error handling when DAO raises an exception.""" + mock_list.side_effect = ToolError("CSS template error") + + async with Client(mcp_server) as client: + request = ListCssTemplatesRequest(page=1, page_size=10) + with pytest.raises(ToolError) as excinfo: # noqa: PT012 + await client.call_tool( + "list_css_templates", {"request": request.model_dump()} + ) + assert "CSS template error" in str(excinfo.value) + + +@patch("superset.daos.css.CssTemplateDAO.find_by_id") [email protected] +async def test_get_css_template_info_basic(mock_find, mcp_server): + """Test basic get CSS template info functionality.""" + template = create_mock_css_template() + mock_find.return_value = template + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_css_template_info", {"request": {"identifier": 1}} + ) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["id"] == 1 + assert data["template_name"] == "my_template" + assert data["uuid"] == "test-css-template-uuid-1" + assert data["css"] == "body { color: red; }" Review Comment: Fixed in 180c85afc3. Added `test_get_css_template_info_by_uuid` that calls the tool with a UUID string and asserts `CssTemplateDAO.find_by_id` is called with `id_column="uuid"`, verifying the UUID resolution path end-to-end. ########## tests/unit_tests/mcp_service/theme/tool/test_theme_tools.py: ########## @@ -0,0 +1,235 @@ +# 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. + +import logging +from unittest.mock import MagicMock, patch + +import pytest +from fastmcp import Client +from fastmcp.exceptions import ToolError +from pydantic import ValidationError + +from superset.mcp_service.app import mcp +from superset.mcp_service.theme.schemas import ( + ListThemesRequest, + ThemeFilter, +) +from superset.utils import json + +logging.basicConfig(level=logging.DEBUG) +logger = logging.getLogger(__name__) + + +class TestThemeFilterSchema: + """Tests for ThemeFilter schema — filterable columns.""" + + def test_invalid_filter_column_rejected(self): + """Columns not in the Literal set must be rejected.""" + with pytest.raises(ValidationError): + ThemeFilter(col="not_a_real_column", opr="eq", value="test") + + def test_valid_filter_column_accepted(self): + """theme_name is a valid filter column.""" + f = ThemeFilter(col="theme_name", opr="eq", value="my_theme") + assert f.col == "theme_name" + + def test_json_data_column_not_filterable(self): + """json_data is not a public filter column.""" + with pytest.raises(ValidationError): + ThemeFilter(col="json_data", opr="eq", value="{}") + + +def create_mock_theme( + theme_id: int = 1, + theme_name: str = "light_theme", + json_data: str = '{"primaryColor": "#1890ff"}', + uuid: str = "test-theme-uuid-1", +) -> MagicMock: + """Factory function to create mock theme objects.""" + theme = MagicMock() + theme.id = theme_id + theme.theme_name = theme_name + theme.json_data = json_data + theme.uuid = uuid + theme.is_system = False + theme.is_system_default = False + theme.is_system_dark = False + theme.changed_on = None + theme.created_on = None + return theme + + [email protected] +def mcp_server(): + return mcp + + [email protected](autouse=True) +def mock_auth(): + """Mock authentication for all tests.""" + from unittest.mock import Mock, patch + + with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user: + mock_user = Mock() + mock_user.id = 1 + mock_user.username = "admin" + mock_get_user.return_value = mock_user + yield mock_get_user + + +@patch("superset.daos.theme.ThemeDAO.list") [email protected] +async def test_list_themes_basic(mock_list, mcp_server): + """Test basic theme listing functionality.""" + theme = create_mock_theme() + mock_list.return_value = ([theme], 1) + + async with Client(mcp_server) as client: + request = ListThemesRequest(page=1, page_size=10) + result = await client.call_tool( + "list_themes", {"request": request.model_dump()} + ) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["themes"] is not None + assert len(data["themes"]) == 1 + assert data["themes"][0]["id"] == 1 + assert data["themes"][0]["theme_name"] == "light_theme" + assert data["themes"][0]["uuid"] == "test-theme-uuid-1" + + +@patch("superset.daos.theme.ThemeDAO.list") [email protected] +async def test_list_themes_default_columns_include_uuid(mock_list, mcp_server): + """Test that uuid is included in default columns for themes.""" + theme = create_mock_theme() + mock_list.return_value = ([theme], 1) + + async with Client(mcp_server) as client: + result = await client.call_tool("list_themes", {}) + data = json.loads(result.content[0].text) + + assert data["columns_requested"] == ["id", "theme_name", "uuid"] + assert data["themes"][0]["uuid"] == "test-theme-uuid-1" + + +@patch("superset.daos.theme.ThemeDAO.list") [email protected] +async def test_list_themes_with_search(mock_list, mcp_server): + """Test theme listing with search functionality.""" + theme = create_mock_theme(theme_name="dark_theme") + mock_list.return_value = ([theme], 1) + + async with Client(mcp_server) as client: + request = ListThemesRequest(page=1, page_size=10, search="dark") + result = await client.call_tool( + "list_themes", {"request": request.model_dump()} + ) + data = json.loads(result.content[0].text) + assert len(data["themes"]) == 1 + assert data["themes"][0]["theme_name"] == "dark_theme" + + +@patch("superset.daos.theme.ThemeDAO.list") [email protected] +async def test_list_themes_with_filters(mock_list, mcp_server): + """Test theme listing with filters.""" + theme = create_mock_theme(theme_name="custom_theme") + mock_list.return_value = ([theme], 1) + + async with Client(mcp_server) as client: + request = ListThemesRequest( + page=1, + page_size=10, + filters=[{"col": "theme_name", "opr": "eq", "value": "custom_theme"}], + ) + result = await client.call_tool( + "list_themes", {"request": request.model_dump()} + ) + data = json.loads(result.content[0].text) + assert len(data["themes"]) == 1 + + +@patch("superset.daos.theme.ThemeDAO.list") [email protected] +async def test_list_themes_api_error(mock_list, mcp_server): + """Test error handling when DAO raises an exception.""" + mock_list.side_effect = ToolError("Theme error") + + async with Client(mcp_server) as client: + request = ListThemesRequest(page=1, page_size=10) + with pytest.raises(ToolError) as excinfo: # noqa: PT012 + await client.call_tool("list_themes", {"request": request.model_dump()}) + assert "Theme error" in str(excinfo.value) + + +@patch("superset.daos.theme.ThemeDAO.find_by_id") [email protected] +async def test_get_theme_info_basic(mock_find, mcp_server): + """Test basic get theme info functionality.""" + theme = create_mock_theme() + mock_find.return_value = theme + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_theme_info", {"request": {"identifier": 1}} + ) + assert result.content is not None + data = json.loads(result.content[0].text) + assert data["id"] == 1 + assert data["theme_name"] == "light_theme" + assert data["uuid"] == "test-theme-uuid-1" + assert data["json_data"] == '{"primaryColor": "#1890ff"}' + + +@patch("superset.daos.theme.ThemeDAO.find_by_id") [email protected] +async def test_get_theme_info_by_uuid(mock_find, mcp_server): + """Test get theme info by UUID.""" + theme = create_mock_theme(uuid="a1b2c3d4-5678-90ab-cdef-1234567890ab") + mock_find.return_value = theme + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_theme_info", + {"request": {"identifier": "a1b2c3d4-5678-90ab-cdef-1234567890ab"}}, + ) + data = json.loads(result.content[0].text) + assert data["id"] == 1 + assert data["uuid"] == "a1b2c3d4-5678-90ab-cdef-1234567890ab" + Review Comment: Fixed in 180c85afc3. Updated `test_get_theme_info_by_uuid` to add `mock_find.assert_called_once_with("a1b2c3d4-...", id_column="uuid", query_options=None)`. The test now validates that `ModelGetInfoCore._find_object` routes UUID strings through the `id_column="uuid"` code path, not a fallback ID lookup. ########## superset/mcp_service/system/tool/get_schema.py: ########## @@ -144,6 +152,42 @@ def _get_database_schema_core() -> ModelGetSchemaCore[ModelSchemaInfo]: ) +def _get_css_template_schema_core() -> ModelGetSchemaCore[ModelSchemaInfo]: + """Create CSS template schema core with dynamically extracted columns.""" + from superset.daos.css import CssTemplateDAO + + return ModelGetSchemaCore( + model_type="css_template", + dao_class=CssTemplateDAO, + output_schema=ModelSchemaInfo, + select_columns=get_css_template_columns(), + sortable_columns=CSS_TEMPLATE_SORTABLE_COLUMNS, + default_columns=CSS_TEMPLATE_DEFAULT_COLUMNS, + search_columns=CSS_TEMPLATE_SEARCH_COLUMNS, + default_sort="changed_on", + default_sort_direction="desc", + logger=logger, + ) Review Comment: Fixed in 180c85afc3. Added `filter_columns_override` parameter to `ModelGetSchemaCore` — when set, `_get_filter_columns` returns it directly instead of querying the DAO. Defined `CSS_TEMPLATE_FILTER_COLUMNS = {"template_name": ["eq", "sw", "ilike"]}` and wired it into `_get_css_template_schema_core`. The `get_schema` response now advertises exactly the columns `ListCssTemplatesRequest` accepts. Updated the test to assert `len(filter_columns) == 1` without requiring a DAO mock. ########## superset/mcp_service/system/tool/get_schema.py: ########## @@ -144,6 +152,42 @@ def _get_database_schema_core() -> ModelGetSchemaCore[ModelSchemaInfo]: ) +def _get_css_template_schema_core() -> ModelGetSchemaCore[ModelSchemaInfo]: + """Create CSS template schema core with dynamically extracted columns.""" + from superset.daos.css import CssTemplateDAO + + return ModelGetSchemaCore( + model_type="css_template", + dao_class=CssTemplateDAO, + output_schema=ModelSchemaInfo, + select_columns=get_css_template_columns(), + sortable_columns=CSS_TEMPLATE_SORTABLE_COLUMNS, + default_columns=CSS_TEMPLATE_DEFAULT_COLUMNS, + search_columns=CSS_TEMPLATE_SEARCH_COLUMNS, + default_sort="changed_on", + default_sort_direction="desc", + logger=logger, + ) + + +def _get_theme_schema_core() -> ModelGetSchemaCore[ModelSchemaInfo]: + """Create theme schema core with dynamically extracted columns.""" + from superset.daos.theme import ThemeDAO + + return ModelGetSchemaCore( + model_type="theme", + dao_class=ThemeDAO, + output_schema=ModelSchemaInfo, + select_columns=get_theme_columns(), + sortable_columns=THEME_SORTABLE_COLUMNS, + default_columns=THEME_DEFAULT_COLUMNS, + search_columns=THEME_SEARCH_COLUMNS, + default_sort="changed_on", + default_sort_direction="desc", + logger=logger, + ) Review Comment: Fixed in 180c85afc3. Same approach as the CSS template fix: defined `THEME_FILTER_COLUMNS = {"theme_name": ["eq", "sw", "ilike"]}` and passed it as `filter_columns_override` in `_get_theme_schema_core`. The advertised filter schema now matches `ListThemesRequest` exactly. -- 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]
