codeant-ai-for-open-source[bot] commented on code in PR #38414: URL: https://github.com/apache/superset/pull/38414#discussion_r2891978597
########## tests/unit_tests/mcp_service/sql_lab/tool/test_save_sql_query.py: ########## @@ -0,0 +1,443 @@ +# 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 save_sql_query MCP tool schemas and logic. +""" + +import importlib +import sys +import types +from unittest.mock import MagicMock, Mock, patch + +import pytest +from pydantic import ValidationError + +from superset.mcp_service.sql_lab.schemas import ( + SaveSqlQueryRequest, + SaveSqlQueryResponse, +) + + +class TestSaveSqlQueryRequest: + """Test SaveSqlQueryRequest schema validation.""" + + def test_valid_request(self) -> None: + req = SaveSqlQueryRequest( + database_id=1, + label="Revenue Query", + sql="SELECT SUM(revenue) FROM sales", + ) + assert req.database_id == 1 + assert req.label == "Revenue Query" + assert req.sql == "SELECT SUM(revenue) FROM sales" + + def test_with_optional_fields(self) -> None: + req = SaveSqlQueryRequest( + database_id=1, + label="Revenue Query", + sql="SELECT 1", + schema="public", + catalog="main", + description="Sums revenue", + ) + assert req.schema_name == "public" + assert req.catalog == "main" + assert req.description == "Sums revenue" + + def test_empty_sql_fails(self) -> None: + with pytest.raises(ValidationError, match="SQL query cannot be empty"): + SaveSqlQueryRequest(database_id=1, label="test", sql=" ") + + def test_empty_label_fails(self) -> None: + with pytest.raises(ValidationError, match="Label cannot be empty"): + SaveSqlQueryRequest(database_id=1, label=" ", sql="SELECT 1") + + def test_sql_is_stripped(self) -> None: + req = SaveSqlQueryRequest(database_id=1, label="test", sql=" SELECT 1 ") + assert req.sql == "SELECT 1" + + def test_label_is_stripped(self) -> None: + req = SaveSqlQueryRequest(database_id=1, label=" My Query ", sql="SELECT 1") + assert req.label == "My Query" + + def test_label_max_length(self) -> None: + with pytest.raises(ValidationError, match="String should have at most 256"): + SaveSqlQueryRequest(database_id=1, label="a" * 257, sql="SELECT 1") + + def test_schema_alias(self) -> None: + """The field accepts 'schema' as alias for 'schema_name'.""" + req = SaveSqlQueryRequest( + database_id=1, + label="test", + sql="SELECT 1", + schema="public", + ) + assert req.schema_name == "public" + + +class TestSaveSqlQueryResponse: + """Test SaveSqlQueryResponse schema.""" + + def test_response_fields(self) -> None: + resp = SaveSqlQueryResponse( + id=42, + label="Revenue", + sql="SELECT 1", + database_id=1, + url="/sqllab?savedQueryId=42", + ) + assert resp.id == 42 + assert resp.label == "Revenue" + assert resp.url == "/sqllab?savedQueryId=42" + + def test_response_with_optional_fields(self) -> None: + resp = SaveSqlQueryResponse( + id=42, + label="Revenue", + sql="SELECT 1", + database_id=1, + schema="public", + description="A query", + url="/sqllab?savedQueryId=42", + ) + assert resp.schema_name == "public" + assert resp.description == "A query" + + +def _force_passthrough_decorators(): + """Force superset_core.api.mcp.tool to be a passthrough decorator. + + In CI, superset_core is fully installed and the real @tool decorator + includes authentication middleware. For unit tests we want to bypass + auth and test the tool logic directly, so we always replace the + decorator with a passthrough regardless of installation state. + + Returns a dict of original sys.modules entries so they can be restored. + """ + + def _passthrough_tool(func=None, **kwargs): + if func is not None: + return func + return lambda f: f + + mock_mcp = MagicMock() + mock_mcp.tool = _passthrough_tool + + mock_api = MagicMock() + mock_api.mcp = mock_mcp + + # Save original modules so we can restore them later + saved_modules: dict[str, types.ModuleType] = {} + superset_core_keys = [k for k in sys.modules if k.startswith("superset_core")] + for key in superset_core_keys: + saved_modules[key] = sys.modules.pop(key) + + # Mock all possible import paths for superset_core + sys.modules["superset_core"] = MagicMock() + sys.modules["superset_core.api"] = mock_api + sys.modules["superset_core.api.mcp"] = mock_mcp + sys.modules["superset_core.mcp"] = mock_mcp + sys.modules.setdefault("superset_core.api.types", MagicMock()) + + return saved_modules + + +def _restore_modules(saved_modules: dict[str, types.ModuleType]) -> None: + """Restore original sys.modules entries after passthrough mocking.""" + # Remove mock entries + for key in list(sys.modules.keys()): + if key.startswith("superset_core"): + del sys.modules[key] + # Restore originals Review Comment: **Suggestion:** The helper that restores `sys.modules` only removes `superset_core` entries, so when the tool module is first imported under the mocked decorator it remains permanently patched in `sys.modules` for the rest of the test run, causing later tests to see an undecorated (auth-less) version of the tool and introducing cross-test contamination. [logic error] <details> <summary><b>Severity Level:</b> Major ⚠️</summary> ```mdx - ⚠️ Later tests may use auth-bypassing `save_sql_query` implementation. - ⚠️ MCP tool authentication and middleware tests can become unreliable. ``` </details> ```suggestion # Remove mock entries and any tool modules imported under patched decorators for key in list(sys.modules.keys()): if key.startswith("superset_core") or key.startswith( "superset.mcp_service.sql_lab.tool" ): del sys.modules[key] # Restore originals (including any previously-imported tool modules) ``` <details> <summary><b>Steps of Reproduction ✅ </b></summary> ```mdx 1. Run any test in `tests/unit_tests/mcp_service/sql_lab/tool/test_save_sql_query.py` that calls `_get_tool_module()` (e.g. `TestSaveSqlQueryToolLogic.test_save_query_creates_saved_query` at lines ~214–283). `_get_tool_module()` (lines 170–186) calls `_force_passthrough_decorators()` (lines 122–157) to mock `superset_core.*`, then imports `superset.mcp_service.sql_lab.tool.save_sql_query` under the passthrough `@tool` decorator and returns `(mod, saved_modules)`. 2. At this first call, there is no pre-existing `superset.mcp_service.sql_lab.tool.save_sql_query` in `sys.modules`, so `_get_tool_module()`'s loop over `sys.modules` (lines 176–183) does not add any tool modules to `saved_tool_modules`/`saved_modules`. The only entries saved are the original `superset_core*` modules. 3. When the `try` block in the test finishes, the `finally` block calls `_restore_modules(saved_modules)` (lines 160–167). That helper only deletes keys starting with `"superset_core"` (lines 163–165), leaving the patched `superset.mcp_service.sql_lab.tool.save_sql_query` module (imported in step 1) resident in `sys.modules`. The subsequent `sys.modules.update(saved_modules)` (line 167) restores the original `superset_core*` entries but does not overwrite the tool module because it was not present in `saved_modules`. 4. Any later test in the same pytest process that does a plain `import superset.mcp_service.sql_lab.tool.save_sql_query` (without going through `_get_tool_module()`) will now receive the already-imported, auth-bypassing module from `sys.modules` instead of re-importing the original decorated version. This is concrete Python import behavior based on the current `_restore_modules()` implementation and the import path `superset.mcp_service.sql_lab.tool.save_sql_query` used in `_get_tool_module()` (line 175), so cross-test contamination is real rather than hypothetical. ``` </details> <details> <summary><b>Prompt for AI Agent 🤖 </b></summary> ```mdx This is a comment left during a code review. **Path:** tests/unit_tests/mcp_service/sql_lab/tool/test_save_sql_query.py **Line:** 162:166 **Comment:** *Logic Error: The helper that restores `sys.modules` only removes `superset_core` entries, so when the tool module is first imported under the mocked decorator it remains permanently patched in `sys.modules` for the rest of the test run, causing later tests to see an undecorated (auth-less) version of the tool and introducing cross-test contamination. Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise. ``` </details> <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F38414&comment_hash=7e10f6edf905085fc539a02dd11b631aef351db638b016add6dca00e1537083e&reaction=like'>👍</a> | <a href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F38414&comment_hash=7e10f6edf905085fc539a02dd11b631aef351db638b016add6dca00e1537083e&reaction=dislike'>👎</a> -- 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]
