aminghadersohi commented on code in PR #41924:
URL: https://github.com/apache/superset/pull/41924#discussion_r3573746669
##########
tests/unit_tests/mcp_service/semantic_layer/tool/test_list_metrics.py:
##########
@@ -112,18 +115,54 @@ def _access_denied_exc(message: str = "Access denied") ->
SupersetSecurityExcept
)
[email protected]
-async def test_list_metrics_builtin_happy_path(mcp_server: FastMCP) -> None:
- """list_metrics returns builtin metrics when only datasets exist."""
- mock_ds = _make_dataset(42)
[email protected]
+def _patched_dataset_lookup(
+ dataset: MagicMock | None,
+) -> Generator[tuple[MagicMock, MagicMock], None, None]:
+ """Patch the ``dataset_id``-driven lookup path.
+
+ Covers ``DatasetDAO.find_by_id`` (direct lookup by id) plus the
+ ``SemanticViewDAO.find_accessible`` call every ``list_metrics`` request
+ makes regardless of scope.
+ """
+ with (
+ patch.object(list_metrics_module, "DatasetDAO") as mock_dao,
+ patch.object(list_metrics_module, "SemanticViewDAO") as mock_view_dao,
+ ):
+ mock_dao.find_by_id.return_value = dataset
+ mock_view_dao.find_accessible.return_value = []
+ yield mock_dao, mock_view_dao
+
[email protected]
+def _patched_dataset_search(
+ datasets: list[MagicMock],
+) -> Generator[tuple[MagicMock, MagicMock, MagicMock], None, None]:
+ """Patch the search-driven lookup path.
+
+ Covers the unscoped/searched query path: ``DatasetDAO._apply_base_filter``
+ applied to a ``db.session.query(...).options(...)`` chain, as used when no
+ ``dataset_id``/``view_id`` narrows the request.
+ """
with (
patch.object(list_metrics_module, "DatasetDAO") as mock_dao,
patch.object(list_metrics_module, "SemanticViewDAO") as mock_view_dao,
+ patch.object(list_metrics_module, "db") as mock_db,
):
- mock_dao.find_by_id.return_value = mock_ds
mock_view_dao.find_accessible.return_value = []
+ mock_query: MagicMock = MagicMock()
+ mock_db.session.query.return_value.options.return_value = mock_query
+ mock_dao._apply_base_filter.return_value = mock_query
+ mock_query.all.return_value = datasets
+ yield mock_dao, mock_view_dao, mock_db
+
[email protected]()
+async def test_list_metrics_builtin_happy_path(mcp_server: FastMCP) -> None:
+ """list_metrics returns builtin metrics when only datasets exist."""
+ mock_ds = _make_dataset(42)
Review Comment:
Added the explicit `MagicMock` annotation in 3c21d4d2db.
##########
tests/unit_tests/mcp_service/semantic_layer/tool/test_list_metrics.py:
##########
@@ -169,22 +208,12 @@ async def test_list_metrics_privacy_check(mcp_server:
FastMCP) -> None:
assert data["error_type"] == "DataModelMetadataRestricted"
[email protected]
[email protected]()
async def test_list_metrics_search_filter(mcp_server: FastMCP) -> None:
"""list_metrics filters metrics by search term."""
mock_ds: MagicMock = _make_dataset(1)
Review Comment:
This local already has the requested `MagicMock` annotation from 8c643dd92d;
no further change was needed.
##########
tests/unit_tests/mcp_service/sql_lab/tool/test_open_sql_lab_with_context.py:
##########
@@ -337,3 +337,108 @@ def
test_sanitizes_error_and_keeps_empty_url_for_missing_database(self) -> None:
)
finally:
_restore_modules(saved_modules)
+
+ def test_returns_error_for_invalid_nonexistent_database_id(self) -> None:
+ """DatabaseDAO.find_by_id returning None for an unknown ID must produce
+ a structured not-found error rather than a raw crash."""
+ mod, saved_modules = _get_tool_module()
+ try:
+ request = OpenSqlLabRequest(database_id=999999999, sql="SELECT 1")
Review Comment:
Added the explicit `OpenSqlLabRequest` annotation in 3c21d4d2db.
##########
tests/unit_tests/mcp_service/sql_lab/tool/test_open_sql_lab_with_context.py:
##########
@@ -337,3 +337,108 @@ def
test_sanitizes_error_and_keeps_empty_url_for_missing_database(self) -> None:
)
finally:
_restore_modules(saved_modules)
+
+ def test_returns_error_for_invalid_nonexistent_database_id(self) -> None:
+ """DatabaseDAO.find_by_id returning None for an unknown ID must produce
+ a structured not-found error rather than a raw crash."""
+ mod, saved_modules = _get_tool_module()
+ try:
+ request = OpenSqlLabRequest(database_id=999999999, sql="SELECT 1")
+
+ with (
+ patch(
+ "superset.daos.database.DatabaseDAO.find_by_id",
+ return_value=None,
+ ) as mock_find_by_id,
+ patch.object(
+ mod.event_logger, "log_context", return_value=nullcontext()
+ ),
+ ):
+ response = mod.open_sql_lab_with_context(request,
_make_mock_ctx())
+
+ mock_find_by_id.assert_called_once_with(999999999)
+ assert response.url == ""
+ assert response.database_id == 999999999
+ assert response.error == sanitize_for_llm_context(
+ "Database with ID 999999999 not found."
+ " Use list_databases to get valid database IDs.",
+ field_path=("error",),
+ )
+ finally:
+ _restore_modules(saved_modules)
+
+ def test_returns_generic_not_found_error_when_database_access_denied(
+ self,
+ ) -> None:
+ """The tool has no dedicated permission-denied branch: it relies solely
+ on ``DatabaseDAO.find_by_id``, whose base filter (``DatabaseFilter`` in
+ ``superset/databases/filters.py``) scopes query results to databases
the
+ requesting user can access. A database that exists but that the current
+ user lacks access to is filtered out of the query and ``find_by_id``
+ returns ``None`` -- indistinguishable, from this tool's perspective,
+ from a genuinely nonexistent ID. This test locks in that fail-closed,
+ non-leaking behavior: no distinct "access denied" message is emitted
+ that would reveal the database's existence to an unauthorized caller.
+ """
+ mod, saved_modules = _get_tool_module()
+ try:
+ request = OpenSqlLabRequest(
+ database_id=42,
+ schema="restricted_schema",
+ title="Query I cannot access",
+ )
Review Comment:
Added the explicit `OpenSqlLabRequest` annotation in 3c21d4d2db.
##########
tests/unit_tests/mcp_service/sql_lab/tool/test_open_sql_lab_with_context.py:
##########
@@ -337,3 +337,108 @@ def
test_sanitizes_error_and_keeps_empty_url_for_missing_database(self) -> None:
)
finally:
_restore_modules(saved_modules)
+
+ def test_returns_error_for_invalid_nonexistent_database_id(self) -> None:
+ """DatabaseDAO.find_by_id returning None for an unknown ID must produce
+ a structured not-found error rather than a raw crash."""
+ mod, saved_modules = _get_tool_module()
+ try:
+ request = OpenSqlLabRequest(database_id=999999999, sql="SELECT 1")
+
+ with (
+ patch(
+ "superset.daos.database.DatabaseDAO.find_by_id",
+ return_value=None,
+ ) as mock_find_by_id,
+ patch.object(
+ mod.event_logger, "log_context", return_value=nullcontext()
+ ),
+ ):
+ response = mod.open_sql_lab_with_context(request,
_make_mock_ctx())
+
+ mock_find_by_id.assert_called_once_with(999999999)
+ assert response.url == ""
+ assert response.database_id == 999999999
+ assert response.error == sanitize_for_llm_context(
+ "Database with ID 999999999 not found."
+ " Use list_databases to get valid database IDs.",
+ field_path=("error",),
+ )
+ finally:
+ _restore_modules(saved_modules)
+
+ def test_returns_generic_not_found_error_when_database_access_denied(
+ self,
+ ) -> None:
+ """The tool has no dedicated permission-denied branch: it relies solely
+ on ``DatabaseDAO.find_by_id``, whose base filter (``DatabaseFilter`` in
+ ``superset/databases/filters.py``) scopes query results to databases
the
+ requesting user can access. A database that exists but that the current
+ user lacks access to is filtered out of the query and ``find_by_id``
+ returns ``None`` -- indistinguishable, from this tool's perspective,
+ from a genuinely nonexistent ID. This test locks in that fail-closed,
+ non-leaking behavior: no distinct "access denied" message is emitted
+ that would reveal the database's existence to an unauthorized caller.
+ """
+ mod, saved_modules = _get_tool_module()
+ try:
+ request = OpenSqlLabRequest(
+ database_id=42,
+ schema="restricted_schema",
+ title="Query I cannot access",
+ )
+
+ with (
+ patch(
+ "superset.daos.database.DatabaseDAO.find_by_id",
+ return_value=None,
+ ) as mock_find_by_id,
+ patch.object(
+ mod.event_logger, "log_context", return_value=nullcontext()
+ ),
+ ):
+ response = mod.open_sql_lab_with_context(request,
_make_mock_ctx())
+
+ mock_find_by_id.assert_called_once_with(42)
+ assert response.url == ""
+ assert response.database_id == 42
+ assert response.schema_name == "restricted_schema"
+ assert response.error == sanitize_for_llm_context(
+ "Database with ID 42 not found."
+ " Use list_databases to get valid database IDs.",
+ field_path=("error",),
+ )
+ finally:
+ _restore_modules(saved_modules)
+
+ def test_returns_error_and_rolls_back_session_on_unexpected_exception(
+ self,
+ ) -> None:
+ """Any unexpected error during DB validation (e.g. a broken connection)
+ must be caught by the outermost handler, roll back the session, and
+ surface a structured error instead of propagating a raw exception."""
+ mod, saved_modules = _get_tool_module()
+ try:
+ request = OpenSqlLabRequest(database_id=7, sql="SELECT 1")
Review Comment:
Added the explicit `OpenSqlLabRequest` annotation in 3c21d4d2db.
##########
tests/unit_tests/mcp_service/system/tool/test_health_check.py:
##########
@@ -17,7 +17,40 @@
"""Tests for health_check MCP tool."""
+import importlib
+from collections.abc import Iterator
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+from fastmcp import Client, FastMCP
+from flask import Flask
+
+from superset.mcp_service.app import mcp
from superset.mcp_service.system.schemas import HealthCheckResponse
+from superset.utils import json
+
+# Import the submodule directly so ``patch.object`` targets the module (not the
+# ``health_check`` function that ``tool/__init__.py`` re-exports onto the
+# package).
+health_check_module = importlib.import_module(
+ "superset.mcp_service.system.tool.health_check"
+)
Review Comment:
Added the explicit `ModuleType` annotation in 3c21d4d2db.
--
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]
