diff --git a/tests/unit_tests/mcp_service/database/tool/test_database_tools.py b/tests/unit_tests/mcp_service/database/tool/test_database_tools.py index 01243977846..01b0661bb2b 100644 --- a/tests/unit_tests/mcp_service/database/tool/test_database_tools.py +++ b/tests/unit_tests/mcp_service/database/tool/test_database_tools.py @@ -18,14 +18,16 @@ import importlib import logging +from collections.abc import Iterator from unittest.mock import MagicMock, patch import pytest -from fastmcp import Client +from fastmcp import Client, FastMCP from fastmcp.exceptions import ToolError from pydantic import ValidationError from superset.mcp_service.app import mcp +from superset.mcp_service.constants import MAX_PAGE_SIZE from superset.mcp_service.database.schemas import DatabaseFilter, ListDatabasesRequest from superset.mcp_service.privacy import DATA_MODEL_METADATA_ERROR_TYPE from superset.utils import json @@ -120,7 +122,7 @@ def mock_auth(): @pytest.fixture(autouse=True) -def allow_data_model_metadata(): +def _allow_data_model_metadata() -> Iterator[None]: """Keep database tests in the normal metadata-allowed path by default.""" with ( patch.object( @@ -378,3 +380,71 @@ async def test_list_databases_does_not_expose_sensitive_credential_columns( # Verify the exploit path: DAO must never receive sensitive column names. dao_columns = mock_list.call_args.kwargs["columns"] assert not sensitive.intersection(dao_columns) + + +# --------------------------------------------------------------------------- +# Pagination edge cases +# --------------------------------------------------------------------------- + + +class TestListDatabasesRequestPagination: + """Schema-level pagination boundary tests — ``page`` is PositiveInt and + ``page_size`` is constrained to (0, MAX_PAGE_SIZE].""" + + def test_page_zero_rejected(self) -> None: + with pytest.raises(ValidationError, match="greater than 0"): + ListDatabasesRequest(page=0) + + def test_negative_page_rejected(self) -> None: + with pytest.raises(ValidationError, match="greater than 0"): + ListDatabasesRequest(page=-1) + + def test_page_size_zero_rejected(self) -> None: + with pytest.raises(ValidationError, match="greater than 0"): + ListDatabasesRequest(page_size=0) + + def test_page_size_over_max_rejected(self) -> None: + with pytest.raises( + ValidationError, + match=f"less than or equal to {MAX_PAGE_SIZE}", + ): + ListDatabasesRequest(page_size=MAX_PAGE_SIZE + 1) + + def test_page_size_at_max_accepted(self) -> None: + request = ListDatabasesRequest(page_size=MAX_PAGE_SIZE) + assert request.page_size == MAX_PAGE_SIZE + + +@pytest.mark.asyncio +async def test_list_databases_invalid_page_size_surfaces_as_tool_error( + mcp_server: FastMCP, +) -> None: + """page_size=0 is rejected before the tool body runs, surfacing as a + structured ToolError rather than a raw 500.""" + async with Client(mcp_server) as client: + with pytest.raises(ToolError, match="greater than 0"): + await client.call_tool("list_databases", {"request": {"page_size": 0}}) + + +@patch("superset.daos.database.DatabaseDAO.list") +@pytest.mark.asyncio +async def test_list_databases_page_beyond_last_page_returns_empty( + mock_list: MagicMock, mcp_server: FastMCP +) -> None: + """A page far past the last page returns an empty list, not an error.""" + # DAO's offset lands past all rows; total_count still reflects the full set. + mock_list.return_value = ([], 2) + async with Client(mcp_server) as client: + request = ListDatabasesRequest(page=9999, page_size=10) + result = await client.call_tool( + "list_databases", {"request": request.model_dump()} + ) + data = json.loads(result.content[0].text) + + assert data["databases"] == [] + assert data["count"] == 0 + assert data["total_count"] == 2 + assert data["page"] == 9999 + assert data["total_pages"] == 1 + assert data["has_next"] is False + assert data["has_previous"] is True diff --git a/tests/unit_tests/mcp_service/rls/tool/test_rls_tools.py b/tests/unit_tests/mcp_service/rls/tool/test_rls_tools.py index eac79a6ba9e..f084467899c 100644 --- a/tests/unit_tests/mcp_service/rls/tool/test_rls_tools.py +++ b/tests/unit_tests/mcp_service/rls/tool/test_rls_tools.py @@ -19,10 +19,12 @@ import logging from unittest.mock import MagicMock, Mock, patch import pytest -from fastmcp import Client +from fastmcp import Client, FastMCP +from fastmcp.exceptions import ToolError from pydantic import ValidationError from superset.mcp_service.app import mcp +from superset.mcp_service.constants import MAX_PAGE_SIZE from superset.mcp_service.rls.schemas import ListRlsFiltersRequest, RlsColumnFilter from superset.utils import json @@ -272,3 +274,71 @@ async def test_get_rls_filter_info_guest_denied(mcp_server): data = json.loads(result.content[0].text) assert data["error_type"] == "Forbidden" assert "guest" in data["error"].lower() + + +# --------------------------------------------------------------------------- +# Pagination edge cases +# --------------------------------------------------------------------------- + + +class TestListRlsFiltersRequestPagination: + """Schema-level pagination boundary tests — ``page`` is PositiveInt and + ``page_size`` is constrained to (0, MAX_PAGE_SIZE].""" + + def test_page_zero_rejected(self) -> None: + with pytest.raises(ValidationError, match="greater than 0"): + ListRlsFiltersRequest(page=0) + + def test_negative_page_rejected(self) -> None: + with pytest.raises(ValidationError, match="greater than 0"): + ListRlsFiltersRequest(page=-1) + + def test_page_size_zero_rejected(self) -> None: + with pytest.raises(ValidationError, match="greater than 0"): + ListRlsFiltersRequest(page_size=0) + + def test_page_size_over_max_rejected(self) -> None: + with pytest.raises( + ValidationError, + match=f"less than or equal to {MAX_PAGE_SIZE}", + ): + ListRlsFiltersRequest(page_size=MAX_PAGE_SIZE + 1) + + def test_page_size_at_max_accepted(self) -> None: + request = ListRlsFiltersRequest(page_size=MAX_PAGE_SIZE) + assert request.page_size == MAX_PAGE_SIZE + + +@pytest.mark.asyncio +async def test_list_rls_filters_invalid_page_size_surfaces_as_tool_error( + mcp_server: FastMCP, +) -> None: + """page_size=0 is rejected before the tool body runs, surfacing as a + structured ToolError rather than a raw 500.""" + async with Client(mcp_server) as client: + with pytest.raises(ToolError, match="greater than 0"): + await client.call_tool("list_rls_filters", {"request": {"page_size": 0}}) + + +@patch("superset.daos.security.RLSDAO.list") +@pytest.mark.asyncio +async def test_list_rls_filters_page_beyond_last_page_returns_empty( + mock_list: MagicMock, mcp_server: FastMCP +) -> None: + """A page far past the last page returns an empty list, not an error.""" + # DAO's offset lands past all rows; total_count still reflects the full set. + mock_list.return_value = ([], 2) + async with Client(mcp_server) as client: + request = ListRlsFiltersRequest(page=9999, page_size=10) + result = await client.call_tool( + "list_rls_filters", {"request": request.model_dump()} + ) + data = json.loads(result.content[0].text) + + assert data["rls_filters"] == [] + assert data["count"] == 0 + assert data["total_count"] == 2 + assert data["page"] == 9999 + assert data["total_pages"] == 1 + assert data["has_next"] is False + assert data["has_previous"] is True diff --git a/tests/unit_tests/mcp_service/semantic_layer/tool/test_get_compatible_dimensions.py b/tests/unit_tests/mcp_service/semantic_layer/tool/test_get_compatible_dimensions.py index 2e1ffcb77ad..85db6c5c892 100644 --- a/tests/unit_tests/mcp_service/semantic_layer/tool/test_get_compatible_dimensions.py +++ b/tests/unit_tests/mcp_service/semantic_layer/tool/test_get_compatible_dimensions.py @@ -22,6 +22,7 @@ from __future__ import annotations import importlib from collections.abc import Generator from types import ModuleType +from typing import Any from unittest.mock import MagicMock, Mock, patch import pytest @@ -282,3 +283,91 @@ async def test_get_compatible_dimensions_external_not_found( assert data["success"] is False assert data["error_type"] == "NotFound" + + +@pytest.mark.asyncio +async def test_get_compatible_dimensions_builtin_empty_selection( + mcp_server: FastMCP, +) -> None: + """Explicitly empty selected_metrics/selected_dimensions is not an error. + + An empty selection is the natural starting state of a query builder + (nothing picked yet), so it must return the full groupby-enabled column + set rather than a validation failure. + """ + mock_ds: MagicMock = _make_dataset(42) + + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_compatible_dimensions", + { + "request": { + "dataset_id": 42, + "selected_metrics": [], + "selected_dimensions": [], + } + }, + ) + data: dict[str, Any] = json.loads(result.content[0].text) + + assert data["success"] is True + names: set[str] = {d["name"] for d in data["compatible_dimensions"]} + assert names == {"region", "category"} + + +@pytest.mark.asyncio +async def test_get_compatible_dimensions_external_empty_selection( + mcp_server: FastMCP, +) -> None: + """External views handle an explicitly empty selection without error.""" + mock_view: MagicMock = _make_view(5) + mock_view.get_compatible_dimensions = MagicMock(return_value=[]) + + with patch( + "superset.daos.semantic_layer.SemanticViewDAO.find_by_id", + return_value=mock_view, + ): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_compatible_dimensions", + { + "request": { + "view_id": 5, + "selected_metrics": [], + "selected_dimensions": [], + } + }, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is True + assert data["compatible_dimensions"] == [] + mock_view.get_compatible_dimensions.assert_called_once_with([], []) + + +@pytest.mark.asyncio +async def test_get_compatible_dimensions_unicode_unknown_selection_validation_error( + mcp_server: FastMCP, +) -> None: + """Unicode/special-character names in an unknown selection surface cleanly.""" + mock_ds = _make_dataset(42) + + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_compatible_dimensions", + { + "request": { + "dataset_id": 42, + "selected_metrics": ["日本語_metric"], + "selected_dimensions": ["special!chars?"], + } + }, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "ValidationError" + assert "Unknown metric: '日本語_metric'" in data["error"] + assert "Unknown dimension: 'special!chars?'" in data["error"] diff --git a/tests/unit_tests/mcp_service/semantic_layer/tool/test_get_compatible_metrics.py b/tests/unit_tests/mcp_service/semantic_layer/tool/test_get_compatible_metrics.py index 763ff8d0f7e..eb89e58d768 100644 --- a/tests/unit_tests/mcp_service/semantic_layer/tool/test_get_compatible_metrics.py +++ b/tests/unit_tests/mcp_service/semantic_layer/tool/test_get_compatible_metrics.py @@ -22,6 +22,7 @@ from __future__ import annotations import importlib from collections.abc import Generator from types import ModuleType +from typing import Any from unittest.mock import MagicMock, Mock, patch import pytest @@ -278,3 +279,91 @@ async def test_get_compatible_metrics_not_found(mcp_server: FastMCP) -> None: assert data["success"] is False assert data["error_type"] == "NotFound" + + +@pytest.mark.asyncio +async def test_get_compatible_metrics_builtin_empty_selection( + mcp_server: FastMCP, +) -> None: + """Explicitly empty selected_metrics/selected_dimensions is not an error. + + An empty selection is the natural starting state of a query builder + (nothing picked yet), so it must return all dataset metrics rather than + a validation failure. + """ + mock_ds: MagicMock = _make_dataset(42) + + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_compatible_metrics", + { + "request": { + "dataset_id": 42, + "selected_metrics": [], + "selected_dimensions": [], + } + }, + ) + data: dict[str, Any] = json.loads(result.content[0].text) + + assert data["success"] is True + names: set[str] = {m["name"] for m in data["compatible_metrics"]} + assert names == {"count", "revenue"} + + +@pytest.mark.asyncio +async def test_get_compatible_metrics_external_empty_selection( + mcp_server: FastMCP, +) -> None: + """External views handle an explicitly empty selection without error.""" + mock_view: MagicMock = _make_view(5) + mock_view.get_compatible_metrics = MagicMock(return_value=[]) + + with patch( + "superset.daos.semantic_layer.SemanticViewDAO.find_by_id", + return_value=mock_view, + ): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_compatible_metrics", + { + "request": { + "view_id": 5, + "selected_metrics": [], + "selected_dimensions": [], + } + }, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is True + assert data["compatible_metrics"] == [] + mock_view.get_compatible_metrics.assert_called_once_with([], []) + + +@pytest.mark.asyncio +async def test_get_compatible_metrics_unicode_unknown_selection_validation_error( + mcp_server: FastMCP, +) -> None: + """Unicode/special-character names in an unknown selection surface cleanly.""" + mock_ds: MagicMock = _make_dataset(42) + + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_compatible_metrics", + { + "request": { + "dataset_id": 42, + "selected_metrics": ["日本語_metric"], + "selected_dimensions": ["special!chars?"], + } + }, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "ValidationError" + assert "Unknown metric: '日本語_metric'" in data["error"] + assert "Unknown dimension: 'special!chars?'" in data["error"] diff --git a/tests/unit_tests/mcp_service/semantic_layer/tool/test_get_table.py b/tests/unit_tests/mcp_service/semantic_layer/tool/test_get_table.py index bbd4f600eb6..f5725bdc3fa 100644 --- a/tests/unit_tests/mcp_service/semantic_layer/tool/test_get_table.py +++ b/tests/unit_tests/mcp_service/semantic_layer/tool/test_get_table.py @@ -22,6 +22,7 @@ from __future__ import annotations import importlib from collections.abc import Generator from types import ModuleType +from typing import Any from unittest.mock import MagicMock, Mock, patch import pytest @@ -297,3 +298,231 @@ async def test_get_table_external_time_range_without_dttm_validation_error( assert data["success"] is False assert data["error_type"] == "ValidationError" assert "no datetime dimension" in data["message"] + + +@pytest.mark.asyncio +async def test_get_table_dataset_not_found(mcp_server: FastMCP) -> None: + """get_table returns NotFound when dataset_id doesn't resolve to a dataset.""" + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=None): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + {"request": {"dataset_id": 999999, "metrics": ["revenue"]}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "NotFound" + assert "999999" in data["message"] + + +@pytest.mark.asyncio +async def test_get_table_view_not_found(mcp_server: FastMCP) -> None: + """get_table returns NotFound when view_id doesn't resolve to a view.""" + with patch( + "superset.daos.semantic_layer.SemanticViewDAO.find_by_id", + return_value=None, + ): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + {"request": {"view_id": 999999, "metrics": ["bookings"]}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "NotFound" + assert "999999" in data["message"] + + +@pytest.mark.asyncio +async def test_get_table_invalid_filter_column_validation_error( + mcp_server: FastMCP, +) -> None: + """get_table errors when a filter references an unknown column.""" + mock_ds = _make_dataset(42) + + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + { + "request": { + "dataset_id": 42, + "metrics": ["revenue"], + "filters": [{"col": "bogus_col", "op": "==", "val": "x"}], + } + }, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "ValidationError" + assert "Unknown filter column: 'bogus_col'" in data["error"] + + +@pytest.mark.asyncio +async def test_get_table_invalid_order_by_validation_error( + mcp_server: FastMCP, +) -> None: + """get_table errors when order_by references an unknown column/metric.""" + mock_ds = _make_dataset(42) + + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + { + "request": { + "dataset_id": 42, + "metrics": ["revenue"], + "order_by": ["bogus_order_col"], + } + }, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "ValidationError" + assert "Unknown order_by: 'bogus_order_col'" in data["error"] + + +@pytest.mark.asyncio +async def test_get_table_unknown_filter_operator_passes_through( + mcp_server: FastMCP, +) -> None: + """An operator string outside the documented set is not schema-validated. + + ``GetTableFilter.op`` is a plain ``str`` field (not a Literal/Enum), so + the tool does not reject unrecognized operator values itself -- it + forwards them verbatim to the query layer, which is responsible for + interpreting/rejecting them. + """ + mock_ds = _make_dataset(42) + query_result: dict[str, Any] = { + "queries": [ + { + "data": [{"region": "west", "revenue": 100}], + "colnames": ["region", "revenue"], + "rowcount": 1, + } + ] + } + + with ( + patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds), + patch( + "superset.commands.chart.data.get_data_command.ChartDataCommand" + ) as mock_command_cls, + patch( + "superset.common.query_context_factory.QueryContextFactory" + ) as mock_factory_cls, + ): + mock_command_cls.return_value.run.return_value = query_result + mock_factory_cls.return_value.create.return_value = MagicMock() + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + { + "request": { + "dataset_id": 42, + "metrics": ["revenue"], + "filters": [ + {"col": "region", "op": "TOTALLY_BOGUS_OP", "val": "x"} + ], + } + }, + ) + data = json.loads(result.content[0].text) + + create_kwargs: dict[str, Any] = ( + mock_factory_cls.return_value.create.call_args.kwargs + ) + forwarded_filters: list[dict[str, Any]] = create_kwargs["queries"][0]["filters"] + + assert data["success"] is True + assert {"col": "region", "op": "TOTALLY_BOGUS_OP", "val": "x"} in forwarded_filters + + +@pytest.mark.asyncio +async def test_get_table_unicode_filter_value_passes_through( + mcp_server: FastMCP, +) -> None: + """Unicode filter values are forwarded to the query layer unmodified.""" + mock_ds = _make_dataset(42) + query_result: dict[str, Any] = { + "queries": [ + { + "data": [{"region": "west", "revenue": 100}], + "colnames": ["region", "revenue"], + "rowcount": 1, + } + ] + } + unicode_val: str = "日本語 café €" + + with ( + patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds), + patch( + "superset.commands.chart.data.get_data_command.ChartDataCommand" + ) as mock_command_cls, + patch( + "superset.common.query_context_factory.QueryContextFactory" + ) as mock_factory_cls, + ): + mock_command_cls.return_value.run.return_value = query_result + mock_factory_cls.return_value.create.return_value = MagicMock() + + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + { + "request": { + "dataset_id": 42, + "metrics": ["revenue"], + "filters": [{"col": "region", "op": "==", "val": unicode_val}], + } + }, + ) + data = json.loads(result.content[0].text) + + create_kwargs: dict[str, Any] = ( + mock_factory_cls.return_value.create.call_args.kwargs + ) + forwarded_filters: list[dict[str, Any]] = create_kwargs["queries"][0]["filters"] + + assert data["success"] is True + assert {"col": "region", "op": "==", "val": unicode_val} in forwarded_filters + + +@pytest.mark.asyncio +async def test_get_table_builtin_time_range_without_configured_dttm_validation_error( + mcp_server: FastMCP, +) -> None: + """get_table rejects time_range on a builtin dataset with no main_dttm_col. + + Mirrors the external-view "no datetime dimension" case, but for the + builtin path where the datetime column is inferred from + ``dataset.main_dttm_col`` instead of scanning columns. + """ + mock_ds = _make_dataset(42) + mock_ds.main_dttm_col = None + + with patch("superset.daos.dataset.DatasetDAO.find_by_id", return_value=mock_ds): + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_table", + { + "request": { + "dataset_id": 42, + "metrics": ["revenue"], + "time_range": "Last 7 days", + } + }, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is False + assert data["error_type"] == "ValidationError" + assert "no temporal column is configured" in data["message"] diff --git a/tests/unit_tests/mcp_service/semantic_layer/tool/test_list_metrics.py b/tests/unit_tests/mcp_service/semantic_layer/tool/test_list_metrics.py index b083a699112..4777d203a0d 100644 --- a/tests/unit_tests/mcp_service/semantic_layer/tool/test_list_metrics.py +++ b/tests/unit_tests/mcp_service/semantic_layer/tool/test_list_metrics.py @@ -19,13 +19,16 @@ from __future__ import annotations +import contextlib import importlib from collections.abc import Generator from types import ModuleType +from typing import Any from unittest.mock import call, MagicMock, Mock, patch import pytest from fastmcp import Client, FastMCP +from fastmcp.exceptions import ToolError from superset.errors import ErrorLevel, SupersetError, SupersetErrorType from superset.exceptions import SupersetSecurityException @@ -112,18 +115,54 @@ def _access_denied_exc(message: str = "Access denied") -> SupersetSecurityExcept ) -@pytest.mark.asyncio -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) +@contextlib.contextmanager +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 = mock_ds + mock_dao.find_by_id.return_value = dataset mock_view_dao.find_accessible.return_value = [] + yield mock_dao, mock_view_dao + +@contextlib.contextmanager +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_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 + + +@pytest.mark.asyncio +async def test_list_metrics_builtin_happy_path(mcp_server: FastMCP) -> None: + """list_metrics returns builtin metrics when only datasets exist.""" + mock_ds: MagicMock = _make_dataset(42) + + with _patched_dataset_lookup(mock_ds): async with Client(mcp_server) as client: result = await client.call_tool( "list_metrics", @@ -174,17 +213,7 @@ async def test_list_metrics_search_filter(mcp_server: FastMCP) -> None: """list_metrics filters metrics by search term.""" mock_ds: MagicMock = _make_dataset(1) - 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_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 = [mock_ds] - + with _patched_dataset_search([mock_ds]): async with Client(mcp_server) as client: result = await client.call_tool( "list_metrics", @@ -280,23 +309,13 @@ async def test_list_metrics_external_per_metric_compatible_dimensions( @pytest.mark.asyncio async def test_list_metrics_pagination_is_stable(mcp_server: FastMCP) -> None: """Metrics are sorted deterministically before pagination is applied.""" - mock_ds = MagicMock() + mock_ds: MagicMock = MagicMock() mock_ds.id = 1 mock_ds.table_name = "table_1" mock_ds.metrics = [_make_metric("zzz_metric"), _make_metric("aaa_metric")] mock_ds.columns = [] - 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_view_dao.find_accessible.return_value = [] - mock_query = 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 = [mock_ds] - + with _patched_dataset_search([mock_ds]): async with Client(mcp_server) as client: page_1 = await client.call_tool( "list_metrics", {"request": {"page": 1, "page_size": 1}} @@ -309,3 +328,194 @@ async def test_list_metrics_pagination_is_stable(mcp_server: FastMCP) -> None: assert data_1["metrics"][0]["name"] == "aaa_metric" assert data_2["metrics"][0]["name"] == "zzz_metric" + + +@pytest.mark.asyncio +async def test_list_metrics_search_no_match_returns_empty(mcp_server: FastMCP) -> None: + """A search term that matches nothing returns an empty (not error) result.""" + mock_ds: MagicMock = _make_dataset(1) + + with _patched_dataset_search([mock_ds]): + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_metrics", + {"request": {"search": "no_such_metric_anywhere"}}, + ) + data: dict[str, Any] = json.loads(result.content[0].text) + + assert data["success"] is True + assert data["metrics"] == [] + assert data["total_count"] == 0 + + +@pytest.mark.asyncio +async def test_list_metrics_nonexistent_dataset_id_returns_empty( + mcp_server: FastMCP, +) -> None: + """A dataset_id that doesn't resolve to a dataset returns an empty result. + + The tool degrades gracefully (empty list) rather than raising NotFound, + since dataset_id here is a scoping filter, not a required lookup key. + """ + with patch.object(list_metrics_module, "DatasetDAO") as mock_dao: + mock_dao.find_by_id.return_value = None + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_metrics", + {"request": {"dataset_id": 999999}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is True + assert data["metrics"] == [] + assert data["total_count"] == 0 + mock_dao.find_by_id.assert_called_once() + + +@pytest.mark.asyncio +async def test_list_metrics_nonexistent_view_id_returns_empty( + mcp_server: FastMCP, +) -> None: + """A view_id that doesn't resolve to a view returns an empty result.""" + with patch.object(list_metrics_module, "SemanticViewDAO") as mock_view_dao: + mock_view_dao.find_by_id.return_value = None + + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_metrics", + {"request": {"view_id": 999999}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is True + assert data["metrics"] == [] + assert data["total_count"] == 0 + mock_view_dao.find_by_id.assert_called_once() + + +@pytest.mark.asyncio +async def test_list_metrics_search_unicode_matches(mcp_server: FastMCP) -> None: + """Unicode search strings match against unicode descriptions correctly.""" + mock_ds: MagicMock = _make_dataset(1) + mock_ds.metrics[1].description = "café blend revenue – daily" + + with _patched_dataset_search([mock_ds]): + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_metrics", + {"request": {"search": "café"}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is True + metrics: list[dict[str, Any]] = data["metrics"] + assert len(metrics) == 1 + assert metrics[0]["name"] == "revenue" + + +@pytest.mark.asyncio +async def test_list_metrics_search_special_characters_no_crash( + mcp_server: FastMCP, +) -> None: + """Search strings with regex-special characters are treated as plain text.""" + mock_ds = _make_dataset(1) + + with _patched_dataset_search([mock_ds]): + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_metrics", + {"request": {"search": "rev$enue%^&*()[.*]"}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is True + assert data["metrics"] == [] + assert data["total_count"] == 0 + + +# --------------------------------------------------------------------------- +# Pagination edge cases +# +# list_metrics hand-rolls its own pagination (list slicing) instead of using +# ModelListCore, but the request schema still enforces page >= 1 and +# 1 <= page_size <= 500 (superset/mcp_service/semantic_layer/schemas.py). +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_list_metrics_page_zero_rejected(mcp_server: FastMCP) -> None: + """page must be >= 1; page=0 is rejected before the tool body runs.""" + async with Client(mcp_server) as client: + with pytest.raises(ToolError, match="greater than or equal to 1"): + await client.call_tool("list_metrics", {"request": {"page": 0}}) + + +@pytest.mark.asyncio +async def test_list_metrics_negative_page_rejected(mcp_server: FastMCP) -> None: + """Negative page numbers are rejected the same way as page=0.""" + async with Client(mcp_server) as client: + with pytest.raises(ToolError, match="greater than or equal to 1"): + await client.call_tool("list_metrics", {"request": {"page": -1}}) + + +@pytest.mark.asyncio +async def test_list_metrics_page_size_zero_rejected(mcp_server: FastMCP) -> None: + """page_size must be >= 1; page_size=0 is rejected before the tool body + runs, surfacing as a structured ToolError rather than a raw 500.""" + async with Client(mcp_server) as client: + with pytest.raises(ToolError, match="greater than or equal to 1"): + await client.call_tool("list_metrics", {"request": {"page_size": 0}}) + + +@pytest.mark.asyncio +async def test_list_metrics_page_size_over_max_rejected(mcp_server: FastMCP) -> None: + """page_size above the 500 ceiling is rejected, not silently clamped.""" + async with Client(mcp_server) as client: + with pytest.raises(ToolError, match="less than or equal to 500"): + await client.call_tool("list_metrics", {"request": {"page_size": 501}}) + + +@pytest.mark.asyncio +async def test_list_metrics_page_size_at_max_accepted(mcp_server: FastMCP) -> None: + """page_size == 500 (the max) is accepted and echoed back.""" + mock_ds: MagicMock = _make_dataset(42) + + with _patched_dataset_lookup(mock_ds): + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_metrics", + {"request": {"dataset_id": 42, "page_size": 500}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is True + assert data["page_size"] == 500 + assert data["total_count"] == 2 + + +@pytest.mark.asyncio +async def test_list_metrics_page_beyond_last_page_returns_empty( + mcp_server: FastMCP, +) -> None: + """Requesting a page past the end returns an empty page, not an error. + + Unlike the ModelListCore-backed list tools, MetricList has no + has_next/has_previous fields — only metrics, total_count, page, + page_size, and total_pages. + """ + mock_ds: MagicMock = _make_dataset(42) + + with _patched_dataset_lookup(mock_ds): + async with Client(mcp_server) as client: + result = await client.call_tool( + "list_metrics", + {"request": {"dataset_id": 42, "page": 9999, "page_size": 50}}, + ) + data = json.loads(result.content[0].text) + + assert data["success"] is True + assert data["metrics"] == [] + assert data["total_count"] == 2 + assert data["page"] == 9999 + assert data["total_pages"] == 1 diff --git a/tests/unit_tests/mcp_service/sql_lab/tool/test_open_sql_lab_with_context.py b/tests/unit_tests/mcp_service/sql_lab/tool/test_open_sql_lab_with_context.py index de17d736ee3..0d937b30e3e 100644 --- a/tests/unit_tests/mcp_service/sql_lab/tool/test_open_sql_lab_with_context.py +++ b/tests/unit_tests/mcp_service/sql_lab/tool/test_open_sql_lab_with_context.py @@ -105,7 +105,7 @@ class TestOpenSqlLabWithContext: def test_sanitizes_direct_sql_and_title_in_url_and_response(self) -> None: mod, saved_modules = _get_tool_module() try: - request = OpenSqlLabRequest( + request: OpenSqlLabRequest = OpenSqlLabRequest( database_id=7, schema="analytics", sql="SELECT * FROM users LIMIT 10", @@ -159,7 +159,7 @@ class TestOpenSqlLabWithContext: def test_sanitizes_generated_dataset_context_sql(self) -> None: mod, saved_modules = _get_tool_module() try: - request = OpenSqlLabRequest( + request: OpenSqlLabRequest = OpenSqlLabRequest( database_id=12, schema="public", dataset_in_context="orders", @@ -203,7 +203,7 @@ class TestOpenSqlLabWithContext: def test_sanitizes_dataset_context_without_schema(self) -> None: mod, saved_modules = _get_tool_module() try: - request = OpenSqlLabRequest( + request: OpenSqlLabRequest = OpenSqlLabRequest( database_id=12, dataset_in_context="orders", ) @@ -276,7 +276,7 @@ class TestOpenSqlLabWithContext: """Whitespace-only titles must not produce a blank-looking tab label.""" mod, saved_modules = _get_tool_module() try: - request = OpenSqlLabRequest( + request: OpenSqlLabRequest = OpenSqlLabRequest( database_id=7, sql="SELECT 1", title=" ", @@ -307,7 +307,7 @@ class TestOpenSqlLabWithContext: def test_sanitizes_error_and_keeps_empty_url_for_missing_database(self) -> None: mod, saved_modules = _get_tool_module() try: - request = OpenSqlLabRequest( + request: OpenSqlLabRequest = OpenSqlLabRequest( database_id=404, schema="analytics", title="Missing database", @@ -337,3 +337,112 @@ class TestOpenSqlLabWithContext: ) 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 = 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 = 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 = OpenSqlLabRequest( + database_id=7, sql="SELECT 1" + ) + + with ( + patch( + "superset.daos.database.DatabaseDAO.find_by_id", + side_effect=RuntimeError("connection reset"), + ), + patch.object( + mod.event_logger, "log_context", return_value=nullcontext() + ), + patch.object(mod.db.session, "rollback") as mock_rollback, + ): + response = mod.open_sql_lab_with_context(request, _make_mock_ctx()) + + mock_rollback.assert_called_once() + assert response.url == "" + assert response.database_id == 7 + assert response.error == sanitize_for_llm_context( + "Failed to generate SQL Lab URL: connection reset", + field_path=("error",), + ) + finally: + _restore_modules(saved_modules) diff --git a/tests/unit_tests/mcp_service/system/tool/test_health_check.py b/tests/unit_tests/mcp_service/system/tool/test_health_check.py index 06b93fe4cda..921c46b4584 100644 --- a/tests/unit_tests/mcp_service/system/tool/test_health_check.py +++ b/tests/unit_tests/mcp_service/system/tool/test_health_check.py @@ -17,7 +17,41 @@ """Tests for health_check MCP tool.""" +import importlib +from collections.abc import Iterator +from types import ModuleType +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: ModuleType = importlib.import_module( + "superset.mcp_service.system.tool.health_check" +) + + +@pytest.fixture +def mcp_server() -> FastMCP: + return mcp + + +@pytest.fixture(autouse=True) +def mock_auth() -> Iterator[MagicMock]: + """Mock authentication for all tests.""" + 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 def test_health_check_response_schema(): @@ -53,3 +87,103 @@ def test_health_check_response_with_uptime(): ) assert response.uptime_seconds == 123.45 + + +# --------------------------------------------------------------------------- +# Tool-level tests: health_check via MCP Client +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_health_check_success_via_client(mcp_server: FastMCP) -> None: + """Happy path: tool returns a healthy status with real system info.""" + with patch.object( + health_check_module, + "get_version_metadata", + return_value={"version_string": "4.1.0"}, + ): + async with Client(mcp_server) as client: + result = await client.call_tool("health_check", {}) + + data = json.loads(result.content[0].text) + assert data["status"] == "healthy" + assert data["service"] == "Superset MCP Service" + assert data["version"] == "4.1.0" + assert data["timestamp"] is not None + assert data["uptime_seconds"] is not None + assert data["uptime_seconds"] >= 0 + + +@pytest.mark.asyncio +async def test_health_check_uses_configured_app_name( + mcp_server: FastMCP, app: Flask +) -> None: + """service name is derived from the APP_NAME config, not hardcoded.""" + had_app_name = "APP_NAME" in app.config + original_app_name = app.config.get("APP_NAME") + app.config["APP_NAME"] = "Acme Analytics" + try: + with patch.object( + health_check_module, + "get_version_metadata", + return_value={"version_string": "4.1.0"}, + ): + async with Client(mcp_server) as client: + result = await client.call_tool("health_check", {}) + finally: + if had_app_name: + app.config["APP_NAME"] = original_app_name + else: + app.config.pop("APP_NAME", None) + + data = json.loads(result.content[0].text) + assert data["service"] == "Acme Analytics MCP Service" + + +@pytest.mark.asyncio +async def test_health_check_returns_error_status_when_version_metadata_raises( + mcp_server: FastMCP, +) -> None: + """The except branch returns a degraded response instead of raising. + + get_version_metadata() is called inside the try block; when it raises, + health_check must catch the exception and report status="error" with + version="unknown" rather than propagating the failure to the client. + """ + with patch.object( + health_check_module, + "get_version_metadata", + side_effect=RuntimeError("version metadata unavailable"), + ): + async with Client(mcp_server) as client: + result = await client.call_tool("health_check", {}) + + data = json.loads(result.content[0].text) + assert data["status"] == "error" + assert data["version"] == "unknown" + # uptime_seconds is only set on the success path. + assert data["uptime_seconds"] is None + # Fields computed before the try block are unaffected by the failure. + assert data["service"] == "Superset MCP Service" + assert data["python_version"] + assert data["platform"] + assert data["timestamp"] is not None + + +@pytest.mark.asyncio +async def test_health_check_returns_error_status_when_log_context_raises( + mcp_server: FastMCP, +) -> None: + """A failure inside the event_logger.log_context block also degrades gracefully.""" + with patch.object( + health_check_module.event_logger, + "log_context", + side_effect=RuntimeError("logging backend unavailable"), + ): + async with Client(mcp_server) as client: + result = await client.call_tool("health_check", {}) + + data = json.loads(result.content[0].text) + assert data["status"] == "error" + assert data["version"] == "unknown" + assert data["uptime_seconds"] is None diff --git a/tests/unit_tests/mcp_service/tag/tool/test_tag_tools.py b/tests/unit_tests/mcp_service/tag/tool/test_tag_tools.py index 1ccd4fa83c0..167e6380513 100644 --- a/tests/unit_tests/mcp_service/tag/tool/test_tag_tools.py +++ b/tests/unit_tests/mcp_service/tag/tool/test_tag_tools.py @@ -19,11 +19,12 @@ import logging from unittest.mock import MagicMock, patch import pytest -from fastmcp import Client +from fastmcp import Client, FastMCP from fastmcp.exceptions import ToolError from pydantic import ValidationError from superset.mcp_service.app import mcp +from superset.mcp_service.constants import MAX_PAGE_SIZE from superset.mcp_service.tag.schemas import ListTagsRequest, TagFilter from superset.utils import json @@ -266,3 +267,69 @@ async def test_list_tags_default_columns_are_id_name_type(mock_list, mcp_server) assert "type" in tag_obj assert "description" not in tag_obj assert "changed_on" not in tag_obj + + +# --------------------------------------------------------------------------- +# Pagination edge cases +# --------------------------------------------------------------------------- + + +class TestListTagsRequestPagination: + """Schema-level pagination boundary tests — ``page`` is PositiveInt and + ``page_size`` is constrained to (0, MAX_PAGE_SIZE].""" + + def test_page_zero_rejected(self) -> None: + with pytest.raises(ValidationError, match="greater than 0"): + ListTagsRequest(page=0) + + def test_negative_page_rejected(self) -> None: + with pytest.raises(ValidationError, match="greater than 0"): + ListTagsRequest(page=-1) + + def test_page_size_zero_rejected(self) -> None: + with pytest.raises(ValidationError, match="greater than 0"): + ListTagsRequest(page_size=0) + + def test_page_size_over_max_rejected(self) -> None: + with pytest.raises( + ValidationError, + match=f"less than or equal to {MAX_PAGE_SIZE}", + ): + ListTagsRequest(page_size=MAX_PAGE_SIZE + 1) + + def test_page_size_at_max_accepted(self) -> None: + request = ListTagsRequest(page_size=MAX_PAGE_SIZE) + assert request.page_size == MAX_PAGE_SIZE + + +@pytest.mark.asyncio +async def test_list_tags_invalid_page_size_surfaces_as_tool_error( + mcp_server: FastMCP, +) -> None: + """page_size=0 is rejected before the tool body runs, surfacing as a + structured ToolError rather than a raw 500.""" + async with Client(mcp_server) as client: + with pytest.raises(ToolError, match="greater than 0"): + await client.call_tool("list_tags", {"request": {"page_size": 0}}) + + +@patch("superset.daos.tag.TagDAO.list") +@pytest.mark.asyncio +async def test_list_tags_page_beyond_last_page_returns_empty( + mock_list: MagicMock, mcp_server: FastMCP +) -> None: + """A page far past the last page returns an empty list, not an error.""" + # DAO's offset lands past all rows; total_count still reflects the full set. + mock_list.return_value = ([], 3) + async with Client(mcp_server) as client: + request = ListTagsRequest(page=9999, page_size=10) + result = await client.call_tool("list_tags", {"request": request.model_dump()}) + data = json.loads(result.content[0].text) + + assert data["tags"] == [] + assert data["count"] == 0 + assert data["total_count"] == 3 + assert data["page"] == 9999 + assert data["total_pages"] == 1 + assert data["has_next"] is False + assert data["has_previous"] is True diff --git a/tests/unit_tests/mcp_service/test_mcp_e2e_smoke.py b/tests/unit_tests/mcp_service/test_mcp_e2e_smoke.py new file mode 100644 index 00000000000..bad7db9a32b --- /dev/null +++ b/tests/unit_tests/mcp_service/test_mcp_e2e_smoke.py @@ -0,0 +1,228 @@ +# 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. +"""End-to-end smoke test for the real MCP ASGI/HTTP stack. + +Every other test under ``tests/unit_tests/mcp_service/`` drives tools via +FastMCP's in-process ``Client(mcp)``, which talks to the bare ``FastMCP`` +object through an in-memory transport (``FastMCPTransport``). That path +never serializes a JSON-RPC message over HTTP and never runs through the +Starlette ASGI app that ``superset.mcp_service.server.run_server()`` builds +via ``mcp_instance.http_app(...)`` -- so the FastMCP-level middleware stack +(``LoggingMiddleware``, ``GlobalErrorHandlerMiddleware``, +``StructuredContentStripperMiddleware``, etc., see +``build_middleware_list()`` in ``server.py``) is never actually exercised in +CI. + +This module closes that gap with a single smoke test file that: + +1. Builds the *real* ASGI app the way ``run_server()`` does -- + ``mcp.http_app(transport="streamable-http", stateless_http=True)`` -- + with the production FastMCP-level middleware list attached. +2. Serves it in-process over real MCP streamable-HTTP JSON-RPC using + ``httpx.ASGITransport`` (no real TCP socket, no real network). +3. Drives it with FastMCP's own high-level ``Client``, proving the full + request/response wire protocol (headers, session negotiation, JSON-RPC + envelopes) round-trips correctly through the real transport. + +Deliberately out of scope: the Starlette-level ``BrowserHelloMiddleware`` +(added via ``_build_starlette_middleware()`` in ``server.py``) only +intercepts ``GET``/``HEAD`` requests carrying a browser ``Accept`` header +(see ``BrowserHelloMiddleware.dispatch`` in ``jwt_verifier.py``); it never +touches the ``POST /mcp`` JSON-RPC path this test exercises, and wiring it +up would require standing up real auth/config machinery for no additional +coverage here. + +Global state note: the FastMCP middleware list lives on the *shared* +``superset.mcp_service.app.mcp`` singleton that every other test file also +imports (``mcp.middleware`` is a plain list mutated in place by +``add_middleware()``). The ``_real_asgi_client`` helper below snapshots and +restores that list so this file cannot leak middleware into other tests. +""" + +import contextlib +from collections.abc import AsyncIterator, Iterator, Mapping +from typing import Any +from unittest.mock import Mock, patch + +import anyio +import httpx +import pytest +from fastmcp import Client +from fastmcp.client.transports import StreamableHttpTransport +from starlette.applications import Starlette + +from superset.mcp_service.app import mcp +from superset.mcp_service.server import build_middleware_list +from superset.utils import json + +ASGIMessage = Mapping[str, Any] + + +@pytest.fixture(autouse=True) +def mock_auth() -> Iterator[Mock]: + """Make authentication deterministic for list and call requests.""" + 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 + + +@contextlib.asynccontextmanager +async def _run_asgi_lifespan(app: Starlette) -> AsyncIterator[None]: + """Manually drive the ASGI lifespan protocol for ``app``. + + FastMCP's streamable-HTTP app only creates its + ``StreamableHTTPSessionManager`` (and starts its task group) inside the + app's ``lifespan`` context manager -- see ``create_streamable_http_app`` + in ``fastmcp.server.http``. ``httpx.ASGITransport`` drives the ASGI + ``http`` scope but never sends ``lifespan`` events, so without this + helper every request would fail because the session manager was never + started. This is the same protocol libraries like ``asgi-lifespan`` + implement; it's inlined here to avoid a new test-only dependency. + """ + send_stream, receive_stream = anyio.create_memory_object_stream[ASGIMessage](4) + startup_complete = anyio.Event() + shutdown_complete = anyio.Event() + startup_failure: list[str] = [] + shutdown_failure: list[str] = [] + + async def receive() -> ASGIMessage: + return await receive_stream.receive() + + async def send(message: ASGIMessage) -> None: + if message["type"] == "lifespan.startup.complete": + startup_complete.set() + elif message["type"] == "lifespan.startup.failed": + startup_failure.append(message.get("message", "startup failed")) + startup_complete.set() + elif message["type"] == "lifespan.shutdown.complete": + shutdown_complete.set() + elif message["type"] == "lifespan.shutdown.failed": + shutdown_failure.append(message.get("message", "shutdown failed")) + shutdown_complete.set() + + async with anyio.create_task_group() as task_group: + task_group.start_soon( + app, {"type": "lifespan", "asgi": {"version": "3.0"}}, receive, send + ) + await send_stream.send({"type": "lifespan.startup"}) + await startup_complete.wait() + if startup_failure: + raise RuntimeError(f"ASGI app failed to start: {startup_failure[0]}") + try: + yield + finally: + await send_stream.send({"type": "lifespan.shutdown"}) + await shutdown_complete.wait() + if shutdown_failure: + raise RuntimeError( + f"ASGI app failed to shut down: {shutdown_failure[0]}" + ) + + +@contextlib.asynccontextmanager +async def _real_asgi_client() -> AsyncIterator[Client]: + """A FastMCP ``Client`` wired to the real ASGI app over real HTTP semantics. + + Builds the app the way ``run_server()`` does for the multi-pod/http_app + path (``server.py:938``): FastMCP-level middleware from + ``build_middleware_list()`` attached to the shared ``mcp`` instance, then + ``mcp.http_app(transport="streamable-http", stateless_http=True)``. + + The request/response cycle is driven over ``httpx.ASGITransport`` (no + real socket) using FastMCP's own ``StreamableHttpTransport`` so the + client speaks genuine MCP streamable-HTTP JSON-RPC -- initialize + handshake, session headers, and message envelopes -- rather than the + in-process ``FastMCPTransport`` every other test in this package uses. + + Deliberately a plain ``@asynccontextmanager`` used directly by each test + (``async with _real_asgi_client() as client:``) rather than a + yield-based pytest fixture: the anyio task group started inside + ``_run_asgi_lifespan`` enforces that its cancel scope is entered and + exited from the *same* asyncio Task, and pytest-asyncio does not + guarantee that a fixture's pre-yield and post-yield halves run in the + same Task. Keeping setup and teardown inside the test function's own + task sidesteps that entirely. + """ + original_middleware = list(mcp.middleware) + for middleware in build_middleware_list(): + mcp.add_middleware(middleware) + + try: + asgi_app = mcp.http_app(transport="streamable-http", stateless_http=True) + + def httpx_client_factory(**kwargs: Any) -> httpx.AsyncClient: + return httpx.AsyncClient( + transport=httpx.ASGITransport(app=asgi_app), + base_url="http://testserver", + **kwargs, + ) + + transport = StreamableHttpTransport( + "http://testserver/mcp", httpx_client_factory=httpx_client_factory + ) + + async with _run_asgi_lifespan(asgi_app): + async with Client(transport) as client: + yield client + finally: + # Restore the shared FastMCP singleton's middleware list in place + # (not by reassignment) so this file cannot leak state into other + # mcp_service tests, even if something else holds a reference to + # the original list object. + mcp.middleware[:] = original_middleware + + +@pytest.mark.asyncio +async def test_tools_list_over_real_asgi_transport() -> None: + """``tools/list`` round-trips over the real ASGI app + JSON-RPC wire protocol. + + This alone proves the full stack boots: the Starlette app built by + ``http_app()``, the FastMCP-level middleware chain (Logging, + GlobalErrorHandler, StructuredContentStripper, RBAC visibility), the + streamable-HTTP session manager, and real JSON-RPC (de)serialization -- + none of which the in-process ``Client(mcp)`` tests elsewhere in this + package exercise. + """ + async with _real_asgi_client() as client: + tools = await client.list_tools() + + assert len(tools) > 0 + tool_names = {tool.name for tool in tools} + assert "health_check" in tool_names + + +@pytest.mark.asyncio +async def test_tools_call_health_check_over_real_asgi_transport() -> None: + """A real ``tools/call`` for ``health_check`` over the real ASGI transport. + + ``health_check`` is ``@tool(protect=True)`` by default (auth wraps every + tool unless ``protect=False`` is explicit), so authentication is mocked + the same way every other MCP tool test does it: patching + ``get_user_from_request`` via the module-level autouse fixture. That function + is called directly from ``_setup_user_context()`` regardless of how the + request arrived (in-process client vs. real HTTP transport), so the same + mock works unmodified here. + """ + async with _real_asgi_client() as client: + result = await client.call_tool("health_check", {}) + + data = json.loads(result.content[0].text) + assert data["status"] == "healthy" + assert data["service"] == "Superset MCP Service" diff --git a/tests/unit_tests/mcp_service/test_rbac_tool_enforcement.py b/tests/unit_tests/mcp_service/test_rbac_tool_enforcement.py new file mode 100644 index 00000000000..a816ac3028c --- /dev/null +++ b/tests/unit_tests/mcp_service/test_rbac_tool_enforcement.py @@ -0,0 +1,256 @@ +# 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. + +""" +End-to-end RBAC-rejection coverage for mutating MCP tools. + +``tests/unit_tests/mcp_service/conftest.py`` disables RBAC for every other +test module in this package (``disable_mcp_rbac``, autouse), and +``test_auth_rbac.py`` unit-tests ``check_tool_permission`` in isolation by +calling it directly with hand-built stub functions. Neither proves that a +*real, registered* MCP tool actually enforces RBAC when invoked through the +full FastMCP tool-call path (``Client(mcp).call_tool(...)`` -> +``mcp_auth_hook`` -> ``check_tool_permission`` -> the tool body). + +This module closes that gap: for each mutating tool below, it re-enables +RBAC, denies the mocked ``security_manager.can_access`` check, calls the +*actual registered tool* through ``fastmcp.Client``, and asserts the call is +rejected with a ``fastmcp.exceptions.ToolError`` before the tool body runs. +It also includes one control test proving a permitted caller is NOT blocked +by the RBAC gate, so the denial tests above cannot be a false negative caused +by the test harness itself (e.g. a client/transport error that looks like a +rejection for unrelated reasons). + +Scope: only the RBAC-rejection path. General tool behavior (happy path, +validation errors, DAO error handling, etc.) is already covered by each +tool's own test module under ``tests/unit_tests/mcp_service//tool/`` +and is intentionally not duplicated here. +""" + +from collections.abc import Iterator +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from fastmcp import Client +from fastmcp.exceptions import ToolError + +from superset.mcp_service.app import mcp + +# (tool_name, minimal-but-valid request payload, method_permission_name, +# class_permission_name) for each mutating tool audited for RBAC enforcement. +# The request payloads are the smallest bodies that pass each tool's Pydantic +# schema validation (which FastMCP performs while binding arguments, before +# ``mcp_auth_hook`` runs) so the call reaches the RBAC gate itself. +_MUTATING_TOOLS: list[tuple[str, dict[str, Any], str, str]] = [ + ( + "execute_sql", + {"database_id": 1, "sql": "SELECT 1"}, + "execute_sql_query", + "SQLLab", + ), + ( + "update_chart", + {"identifier": 1}, + "write", + "Chart", + ), + ( + "update_dashboard", + {"identifier": 1}, + "write", + "Dashboard", + ), + ( + "generate_chart", + { + "dataset_id": 1, + "config": {"chart_type": "table", "columns": [{"name": "col1"}]}, + }, + "write", + "Chart", + ), + ( + "generate_dashboard", + {"chart_ids": [1]}, + "write", + "Dashboard", + ), + ( + "manage_native_filters", + # ``reorder: []`` satisfies ManageNativeFiltersRequest's "at least one + # operation" validator (checked via ``is None``, not falsiness). + {"dashboard_id": 1, "reorder": []}, + "write", + "Dashboard", + ), + ( + "create_dataset", + {"database_id": 1, "table_name": "my_table"}, + "write", + "Dataset", + ), + ( + "create_virtual_dataset", + {"database_id": 1, "sql": "SELECT 1", "dataset_name": "my_virtual_dataset"}, + "write", + "Dataset", + ), + ( + "create_theme", + {"theme_name": "Denied Theme", "json_data": {"token": {}}}, + "write", + "Theme", + ), + ( + "save_sql_query", + {"database_id": 1, "label": "my query", "sql": "SELECT 1"}, + "write", + "SavedQuery", + ), +] + +_TOOL_IDS: list[str] = [name for name, *_ in _MUTATING_TOOLS] + + +@pytest.fixture +def mcp_server() -> object: + return mcp + + +@pytest.fixture(autouse=True) +def mock_auth() -> Iterator[MagicMock]: + """Authenticate every call as a real (but unprivileged) user. + + Attributes are set explicitly rather than relying on MagicMock's + auto-attribute generation, since ``mcp_auth_hook``/``_setup_user_context`` + reads ``username``, ``id``, ``is_active``, ``roles``, and ``groups`` off + ``g.user``. + """ + with patch("superset.mcp_service.auth.get_user_from_request") as mock_get_user: + mock_user = MagicMock() + mock_user.username = "unprivileged_user" + mock_user.id = 999 + mock_user.is_active = True + mock_user.roles = [] + mock_user.groups = [] + mock_get_user.return_value = mock_user + yield mock_get_user + + +async def _assert_tool_rejected_without_permission( + mcp_server: object, + app: Any, + tool_name: str, + request_payload: dict[str, Any], + method_permission_name: str, + class_permission_name: str, +) -> None: + """Call ``tool_name`` through the real tool-call path with RBAC enabled + and ``security_manager.can_access`` mocked to deny, and assert the call + is rejected by the RBAC gate (``mcp_auth_hook`` -> ``ToolError``) rather + than reaching the tool body.""" + app.config["MCP_RBAC_ENABLED"] = True + try: + mock_sm = MagicMock() + mock_sm.can_access = MagicMock(return_value=False) + with patch("superset.mcp_service.auth.security_manager", mock_sm): + async with Client(mcp_server) as client: + with pytest.raises(ToolError, match="Permission denied"): + await client.call_tool(tool_name, {"request": request_payload}) + mock_sm.can_access.assert_called_with( + f"can_{method_permission_name}", class_permission_name + ) + finally: + app.config.pop("MCP_RBAC_ENABLED", None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tool_name", "request_payload", "method_permission_name", "class_permission_name"), + _MUTATING_TOOLS, + ids=_TOOL_IDS, +) +async def test_mutating_tool_rejects_unauthorized_caller( + mcp_server: object, + app: Any, + tool_name: str, + request_payload: dict[str, Any], + method_permission_name: str, + class_permission_name: str, +) -> None: + """A caller lacking the tool's required FAB permission is rejected by the + real ``mcp_auth_hook`` gate before the tool body runs, for every mutating + tool in the MCP service.""" + await _assert_tool_rejected_without_permission( + mcp_server, + app, + tool_name, + request_payload, + method_permission_name, + class_permission_name, + ) + + +@pytest.mark.asyncio +async def test_authorized_caller_is_not_blocked_by_rbac_gate( + mcp_server: object, app: Any +) -> None: + """Control test: proves the RBAC gate is the thing doing the rejecting + above, not a false negative from the test harness itself. + + Grants ``can_access`` for ``create_theme`` (RBAC enabled) and mocks the + downstream DAO/commit so the call fully succeeds. If the denial tests + above were actually failing for an unrelated reason (bad payload, broken + Client wiring, etc.), this test would fail the same way instead of + succeeding. + """ + import importlib + + create_theme_module = importlib.import_module( + "superset.mcp_service.theme.tool.create_theme" + ) + + mock_theme = MagicMock() + mock_theme.id = 1 + mock_theme.uuid = "11111111-1111-1111-1111-111111111111" + mock_theme.theme_name = "Allowed Theme" + + app.config["MCP_RBAC_ENABLED"] = True + try: + mock_sm = MagicMock() + mock_sm.can_access = MagicMock(return_value=True) + with ( + patch("superset.mcp_service.auth.security_manager", mock_sm), + patch.object(create_theme_module.db.session, "commit"), + patch("superset.daos.theme.ThemeDAO.create", return_value=mock_theme), + ): + async with Client(mcp_server) as client: + result = await client.call_tool( + "create_theme", + { + "request": { + "theme_name": "Allowed Theme", + "json_data": {"token": {"colorPrimary": "#1d4ed8"}}, + } + }, + ) + mock_sm.can_access.assert_called_with("can_write", "Theme") + assert result.structured_content["success"] is True + assert result.structured_content["id"] == 1 + finally: + app.config.pop("MCP_RBAC_ENABLED", None) diff --git a/tests/unit_tests/mcp_service/theme/tool/test_get_theme_info.py b/tests/unit_tests/mcp_service/theme/tool/test_get_theme_info.py index b28aa251ec8..2c132cd857f 100644 --- a/tests/unit_tests/mcp_service/theme/tool/test_get_theme_info.py +++ b/tests/unit_tests/mcp_service/theme/tool/test_get_theme_info.py @@ -131,3 +131,42 @@ async def test_get_theme_info_wraps_json_data( assert "UNTRUSTED-CONTENT" in data["json_data"] assert "fontFamily" in data["json_data"] + + +@patch("superset.daos.theme.ThemeDAO.find_by_id") +@pytest.mark.asyncio +async def test_get_theme_info_invalid_identifier_not_found( + mock_find: MagicMock, mcp_server: object +) -> None: + """A string identifier that is neither an int nor a UUID cannot resolve. + + Themes do not support slug lookups (``supports_slug=False``), so + ``ModelGetInfoCore._find_object`` falls through every branch and returns + None, which surfaces as a ``not_found`` ThemeError — without ever calling + the DAO. + """ + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_theme_info", {"request": {"identifier": "not-a-real-identifier"}} + ) + data = json.loads(result.content[0].text) + assert data["error_type"] == "not_found" + assert "not-a-real-identifier" in data["error"] + mock_find.assert_not_called() + + +@patch("superset.daos.theme.ThemeDAO.find_by_id") +@pytest.mark.asyncio +async def test_get_theme_info_internal_error( + mock_find: MagicMock, mcp_server: object +) -> None: + """Unexpected DAO exceptions are caught and returned as a ThemeError + with error_type='InternalError', rather than propagating.""" + mock_find.side_effect = RuntimeError("database connection lost") + async with Client(mcp_server) as client: + result = await client.call_tool( + "get_theme_info", {"request": {"identifier": 1}} + ) + data = json.loads(result.content[0].text) + assert data["error_type"] == "InternalError" + assert "database connection lost" in data["error"] diff --git a/tests/unit_tests/mcp_service/theme/tool/test_list_themes.py b/tests/unit_tests/mcp_service/theme/tool/test_list_themes.py index c2498df968d..22485f6a975 100644 --- a/tests/unit_tests/mcp_service/theme/tool/test_list_themes.py +++ b/tests/unit_tests/mcp_service/theme/tool/test_list_themes.py @@ -22,9 +22,11 @@ from unittest.mock import MagicMock, Mock, 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.constants import MAX_PAGE_SIZE from superset.mcp_service.theme.schemas import ListThemesRequest, ThemeFilter from superset.utils import json @@ -126,3 +128,174 @@ async def test_list_themes_multiple(mock_list: MagicMock, mcp_server: object) -> names = {t["theme_name"] for t in data["themes"]} assert any("Light" in n for n in names) assert any("Dark" in n for n in names) + + +class TestListThemesRequestPaginationValidation: + """Schema-level validation for page/page_size — mirrors + ListChartsRequest's pagination validation tests, since ListThemesRequest + uses the identical PositiveInt/Field(gt=0, le=MAX_PAGE_SIZE) constraints.""" + + def test_negative_page_rejected(self) -> None: + """page is a PositiveInt (gt=0), so negative values fail schema + validation before the tool body ever runs.""" + with pytest.raises(ValidationError, match="Input should be greater than 0"): + ListThemesRequest(page=-1) + + def test_zero_page_rejected(self) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): + ListThemesRequest(page=0) + + def test_page_size_zero_rejected(self) -> None: + """page_size has gt=0, so 0 is rejected at the schema level.""" + with pytest.raises(ValidationError, match="Input should be greater than 0"): + ListThemesRequest(page_size=0) + + def test_page_size_negative_rejected(self) -> None: + with pytest.raises(ValidationError, match="Input should be greater than 0"): + ListThemesRequest(page_size=-5) + + def test_page_size_exceeds_max_rejected(self) -> None: + """page_size has le=MAX_PAGE_SIZE, so values over the max are rejected + at the schema level — ModelListCore's min(page_size, MAX_PAGE_SIZE) + clamp is defense-in-depth for callers that bypass schema validation, + not something reachable through the public ListThemesRequest path.""" + with pytest.raises( + ValidationError, + match=f"Input should be less than or equal to {MAX_PAGE_SIZE}", + ): + ListThemesRequest(page_size=MAX_PAGE_SIZE + 1) + + def test_page_size_at_max_accepted(self) -> None: + """page_size == MAX_PAGE_SIZE is the boundary-valid case.""" + request = ListThemesRequest(page_size=MAX_PAGE_SIZE) + assert request.page_size == MAX_PAGE_SIZE + + def test_search_and_filters_mutually_exclusive(self) -> None: + """The model_validator rejects combining search with filters.""" + with pytest.raises( + ValidationError, match="Cannot use both 'search' and 'filters'" + ): + ListThemesRequest( + search="blue", + filters=[ThemeFilter(col="theme_name", opr="ct", value="blue")], + ) + + +class TestListThemesPaginationEdges: + """Pagination edge cases exercised through the actual tool call path, + matching ModelListCore.run_tool's real (0-based internally, 1-based in + the response) pagination arithmetic.""" + + @patch("superset.daos.theme.ThemeDAO.list") + @pytest.mark.asyncio + async def test_page_beyond_last_page_returns_empty( + self, mock_list: MagicMock, mcp_server: object + ) -> None: + """Requesting a page past the last page of results returns an empty + themes list but preserves the real total_count/total_pages, since + those come from a separate DAO count independent of the page slice.""" + # Only 1 theme exists in total, but we ask for page 5 (page_size=10): + # a real DB would return no rows for that offset while total_count + # still reflects the overall match count. + mock_list.return_value = ([], 1) + async with Client(mcp_server) as client: + request = ListThemesRequest(page=5, page_size=10) + result = await client.call_tool( + "list_themes", {"request": request.model_dump()} + ) + data = json.loads(result.content[0].text) + + # page is converted to 0-based (page - 1 = 4) before hitting the DAO. + assert mock_list.call_args.kwargs["page"] == 4 + assert mock_list.call_args.kwargs["page_size"] == 10 + assert data["themes"] == [] + assert data["total_count"] == 1 + assert data["total_pages"] == 1 + assert data["page"] == 5 + assert data["has_previous"] is True + assert data["has_next"] is False + + @patch("superset.daos.theme.ThemeDAO.list") + @pytest.mark.asyncio + async def test_page_size_at_max_is_not_clamped( + self, mock_list: MagicMock, mcp_server: object + ) -> None: + """page_size == MAX_PAGE_SIZE passes through untouched (the + defense-in-depth min() clamp in ModelListCore is a no-op here).""" + mock_list.return_value = ([], 0) + async with Client(mcp_server) as client: + request = ListThemesRequest(page=1, page_size=MAX_PAGE_SIZE) + result = await client.call_tool( + "list_themes", {"request": request.model_dump()} + ) + data = json.loads(result.content[0].text) + + assert mock_list.call_args.kwargs["page_size"] == MAX_PAGE_SIZE + assert data["page_size"] == MAX_PAGE_SIZE + + @patch("superset.daos.theme.ThemeDAO.list") + @pytest.mark.asyncio + async def test_empty_result_set( + self, mock_list: MagicMock, mcp_server: object + ) -> None: + """No matching themes: total_pages is 0 (not negative/erroring), and + has_previous/has_next both reflect the first-page, no-results state.""" + mock_list.return_value = ([], 0) + async with Client(mcp_server) as client: + result = await client.call_tool("list_themes", {}) + data = json.loads(result.content[0].text) + + assert data["themes"] == [] + assert data["total_count"] == 0 + assert data["total_pages"] == 0 + assert data["has_previous"] is False + assert data["has_next"] is False + + +@patch("superset.daos.theme.ThemeDAO.list") +@pytest.mark.asyncio +async def test_list_themes_dao_error_propagates_as_tool_error( + mock_list: MagicMock, mcp_server: object +) -> None: + """Unlike get_theme_info, list_themes' outermost except re-raises rather + than returning a ThemeError — the FastMCP middleware converts the + unhandled exception into a ToolError.""" + mock_list.side_effect = RuntimeError("database connection lost") + async with Client(mcp_server) as client: + with pytest.raises(ToolError) as excinfo: # noqa: PT012 + await client.call_tool("list_themes", {}) + assert "database connection lost" in str(excinfo.value) + + +@patch("superset.daos.theme.ThemeDAO.list") +@pytest.mark.asyncio +async def test_list_themes_invalid_select_columns_raises_tool_error( + mock_list: MagicMock, mcp_server: object +) -> None: + """select_columns made up entirely of unknown columns are filtered out + by the ALL_THEME_COLUMNS allowlist in _get_columns_to_load, leaving no + valid columns — ModelListCore raises ValueError, which list_themes lets + propagate (surfaced by middleware as a ToolError).""" + async with Client(mcp_server) as client: + with pytest.raises(ToolError) as excinfo: # noqa: PT012 + await client.call_tool( + "list_themes", {"request": {"select_columns": ["not_a_real_column"]}} + ) + assert "no valid columns" in str(excinfo.value) + mock_list.assert_not_called() + + +@patch("superset.daos.theme.ThemeDAO.list") +@pytest.mark.asyncio +async def test_list_themes_invalid_order_column_raises_tool_error( + mock_list: MagicMock, mcp_server: object +) -> None: + """order_column outside SORTABLE_THEME_COLUMNS is rejected by + _validate_order_column before the DAO is ever queried.""" + async with Client(mcp_server) as client: + with pytest.raises(ToolError) as excinfo: # noqa: PT012 + await client.call_tool( + "list_themes", {"request": {"order_column": "json_data"}} + ) + assert "Invalid order_column" in str(excinfo.value) + mock_list.assert_not_called() diff --git a/tests/unit_tests/mcp_service/user/tool/test_user_tools.py b/tests/unit_tests/mcp_service/user/tool/test_user_tools.py index 92879b2f4e0..877fcaf1c0a 100644 --- a/tests/unit_tests/mcp_service/user/tool/test_user_tools.py +++ b/tests/unit_tests/mcp_service/user/tool/test_user_tools.py @@ -21,11 +21,12 @@ import importlib from unittest.mock import MagicMock, Mock, patch import pytest -from fastmcp import Client +from fastmcp import Client, FastMCP from fastmcp.exceptions import ToolError from pydantic import ValidationError from superset.mcp_service.app import mcp +from superset.mcp_service.constants import MAX_PAGE_SIZE from superset.mcp_service.user.schemas import ListUsersRequest, UserFilter from superset.utils import json @@ -82,7 +83,7 @@ def mock_auth(): @pytest.fixture(autouse=True) -def allow_data_model_metadata(): +def _allow_data_model_metadata(): """Keep user tests in the metadata-allowed path by default.""" with ( patch.object( @@ -119,6 +120,34 @@ class TestUserFilterSchema: assert f.col == "username" +class TestListUsersRequestPagination: + """Schema-level pagination boundary tests — ``page`` is PositiveInt and + ``page_size`` is constrained to (0, MAX_PAGE_SIZE].""" + + def test_page_zero_rejected(self) -> None: + with pytest.raises(ValidationError, match="greater than 0"): + ListUsersRequest(page=0) + + def test_negative_page_rejected(self) -> None: + with pytest.raises(ValidationError, match="greater than 0"): + ListUsersRequest(page=-1) + + def test_page_size_zero_rejected(self) -> None: + with pytest.raises(ValidationError, match="greater than 0"): + ListUsersRequest(page_size=0) + + def test_page_size_over_max_rejected(self) -> None: + with pytest.raises( + ValidationError, + match=f"less than or equal to {MAX_PAGE_SIZE}", + ): + ListUsersRequest(page_size=MAX_PAGE_SIZE + 1) + + def test_page_size_at_max_accepted(self) -> None: + request = ListUsersRequest(page_size=MAX_PAGE_SIZE) + assert request.page_size == MAX_PAGE_SIZE + + # --------------------------------------------------------------------------- # list_users tool tests # --------------------------------------------------------------------------- @@ -290,6 +319,39 @@ async def test_list_users_search_and_filters_mutually_exclusive(mcp_server): ) +@pytest.mark.asyncio +async def test_list_users_invalid_page_size_surfaces_as_tool_error( + mcp_server: FastMCP, +) -> None: + """page_size=0 is rejected before the tool body runs, surfacing as a + structured ToolError rather than a raw 500.""" + async with Client(mcp_server) as client: + with pytest.raises(ToolError, match="greater than 0"): + await client.call_tool("list_users", {"request": {"page_size": 0}}) + + +@patch("superset.daos.user.UserDAO.list") +@pytest.mark.asyncio +async def test_list_users_page_beyond_last_page_returns_empty( + mock_list: MagicMock, mcp_server: FastMCP +) -> None: + """A page far past the last page returns an empty list, not an error.""" + # DAO's offset lands past all rows; total_count still reflects the full set. + mock_list.return_value = ([], 1) + async with Client(mcp_server) as client: + request = ListUsersRequest(page=9999, page_size=10) + result = await client.call_tool("list_users", {"request": request.model_dump()}) + data = json.loads(result.content[0].text) + + assert data["users"] == [] + assert data["count"] == 0 + assert data["total_count"] == 1 + assert data["page"] == 9999 + assert data["total_pages"] == 1 + assert data["has_next"] is False + assert data["has_previous"] is True + + # --------------------------------------------------------------------------- # get_user_info tool tests # --------------------------------------------------------------------------- diff --git a/tests/unit_tests/mcp_service/utils/test_cache_utils.py b/tests/unit_tests/mcp_service/utils/test_cache_utils.py new file mode 100644 index 00000000000..58332ce778f --- /dev/null +++ b/tests/unit_tests/mcp_service/utils/test_cache_utils.py @@ -0,0 +1,271 @@ +# 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 MCP service cache utilities. +""" + +from datetime import datetime, timedelta, timezone +from typing import Any + +from superset.mcp_service.common.cache_schemas import CacheStatus +from superset.mcp_service.utils.cache_utils import ( + apply_cache_control_to_query_context, + get_cache_key_info, + get_cache_status_from_result, + should_use_metadata_cache, +) + +# --------------------------------------------------------------------------- +# get_cache_status_from_result +# --------------------------------------------------------------------------- + + +def test_get_cache_status_from_result_reads_from_queries_list() -> None: + """Should pull cache_hit from the first entry of a 'queries' list.""" + result: dict[str, Any] = {"queries": [{"is_cached": True}]} + + status: CacheStatus = get_cache_status_from_result(result) + + assert status.cache_hit is True + assert status.cache_type == "query" + + +def test_get_cache_status_from_result_falls_back_to_top_level_dict() -> None: + """Should read cache info directly from the result when no 'queries' key.""" + result: dict[str, Any] = {"is_cached": False} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is False + assert status.cache_type == "none" + + +def test_get_cache_status_from_result_empty_queries_list_uses_top_level() -> None: + """Should fall back to the top-level dict when 'queries' is an empty list.""" + result: dict[str, Any] = {"queries": [], "is_cached": True} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is True + + +def test_get_cache_status_from_result_defaults_cache_hit_to_false() -> None: + """Should treat a missing 'is_cached' key as not cached.""" + result: dict[str, Any] = {} + + status = get_cache_status_from_result(result) + + assert status.cache_hit is False + assert status.cache_type == "none" + + +def test_get_cache_status_from_result_reflects_force_refresh_flag() -> None: + """Should surface the caller-provided force_refresh value on the status.""" + result: dict[str, Any] = {"is_cached": True} + + status = get_cache_status_from_result(result, force_refresh=True) + + assert status.refreshed is True + + +def test_get_cache_status_from_result_parses_iso_string_cache_age() -> None: + """Should compute cache_age_seconds from an ISO-formatted cache_dttm string.""" + cache_dt: datetime = datetime.now(timezone.utc) - timedelta(seconds=120) + result: dict[str, Any] = { + "is_cached": True, + "cache_dttm": cache_dt.isoformat(), + } + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is not None + assert status.cache_age_seconds >= 119 + + +def test_get_cache_status_from_result_parses_z_suffixed_string() -> None: + """Should handle a trailing 'Z' UTC designator in the cache_dttm string.""" + result: dict[str, Any] = { + "is_cached": True, + "cache_dttm": "2020-01-01T00:00:00Z", + } + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is not None + assert status.cache_age_seconds > 0 + + +def test_get_cache_status_from_result_parses_datetime_object_cache_age() -> None: + """Should compute cache_age_seconds when cache_dttm is already a datetime.""" + cache_dt = datetime.now(timezone.utc) - timedelta(seconds=60) + result: dict[str, Any] = {"is_cached": True, "cache_dttm": cache_dt} + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is not None + assert status.cache_age_seconds >= 59 + + +def test_get_cache_status_from_result_handles_unparseable_cache_dttm() -> None: + """Should swallow parse errors and leave cache_age_seconds as None.""" + result: dict[str, Any] = {"is_cached": True, "cache_dttm": "not-a-date"} + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is None + # The rest of the status should still be populated correctly. + assert status.cache_hit is True + + +def test_get_cache_status_from_result_no_cache_dttm_leaves_age_none() -> None: + """Should leave cache_age_seconds as None when cache_dttm is absent.""" + result: dict[str, Any] = {"is_cached": True} + + status = get_cache_status_from_result(result) + + assert status.cache_age_seconds is None + + +# --------------------------------------------------------------------------- +# apply_cache_control_to_query_context +# --------------------------------------------------------------------------- + + +def test_apply_cache_control_sets_force_when_cache_disabled() -> None: + """Should set force=True on the query context when use_cache is False.""" + ctx: dict[str, Any] = apply_cache_control_to_query_context( + {"queries": []}, use_cache=False + ) + + assert ctx["force"] is True + + +def test_apply_cache_control_sets_force_when_force_refresh_requested() -> None: + """Should set force=True on the query context when force_refresh is True.""" + ctx = apply_cache_control_to_query_context( + {"queries": []}, use_cache=True, force_refresh=True + ) + + assert ctx["force"] is True + + +def test_apply_cache_control_does_not_set_force_when_using_cache() -> None: + """Should leave 'force' unset when caching is enabled and not refreshing.""" + ctx = apply_cache_control_to_query_context( + {"queries": []}, use_cache=True, force_refresh=False + ) + + assert "force" not in ctx + + +def test_apply_cache_control_applies_cache_timeout_to_every_query() -> None: + """Should stamp cache_timeout onto every query in the context.""" + ctx = apply_cache_control_to_query_context( + {"queries": [{"metric": "a"}, {"metric": "b"}]}, cache_timeout=60 + ) + + assert ctx["queries"][0]["cache_timeout"] == 60 + assert ctx["queries"][1]["cache_timeout"] == 60 + + +def test_apply_cache_control_leaves_queries_untouched_when_timeout_is_none() -> None: + """Should not add a cache_timeout key when none is provided.""" + ctx = apply_cache_control_to_query_context( + {"queries": [{"metric": "a"}]}, cache_timeout=None + ) + + assert "cache_timeout" not in ctx["queries"][0] + + +def test_apply_cache_control_missing_queries_key_is_a_noop_for_timeout() -> None: + """Should not raise or add a 'queries' key when applying a cache_timeout.""" + ctx = apply_cache_control_to_query_context({}, cache_timeout=60) + + assert "queries" not in ctx + + +def test_apply_cache_control_returns_the_same_dict_instance() -> None: + """Should mutate and return the same query_context object it was given.""" + original: dict[str, Any] = {"queries": []} + + result = apply_cache_control_to_query_context(original, use_cache=False) + + assert result is original + + +# --------------------------------------------------------------------------- +# should_use_metadata_cache +# --------------------------------------------------------------------------- + + +def test_should_use_metadata_cache_true_by_default() -> None: + """Should default to using the metadata cache.""" + assert should_use_metadata_cache() is True + + +def test_should_use_metadata_cache_false_when_cache_disabled() -> None: + """Should return False when use_cache is False.""" + assert should_use_metadata_cache(use_cache=False) is False + + +def test_should_use_metadata_cache_false_when_refresh_requested() -> None: + """Should return False when refresh_metadata is True, even with use_cache.""" + assert should_use_metadata_cache(use_cache=True, refresh_metadata=True) is False + + +def test_should_use_metadata_cache_true_when_enabled_and_not_refreshing() -> None: + """Should return True when caching is enabled and no refresh is requested.""" + assert should_use_metadata_cache(use_cache=True, refresh_metadata=False) is True + + +# --------------------------------------------------------------------------- +# get_cache_key_info +# --------------------------------------------------------------------------- + + +def test_get_cache_key_info_returns_none_for_none_input() -> None: + """Should return None when no cache key is provided.""" + assert get_cache_key_info(None) is None + + +def test_get_cache_key_info_returns_none_for_empty_string() -> None: + """Should treat an empty string cache key the same as no key.""" + assert get_cache_key_info("") is None + + +def test_get_cache_key_info_returns_short_keys_unchanged() -> None: + """Should return short cache keys unmodified.""" + key: str = "short_cache_key" + assert get_cache_key_info(key) == key + + +def test_get_cache_key_info_truncates_long_keys() -> None: + """Should truncate cache keys longer than 50 characters with an ellipsis.""" + key = "a" * 60 + + result = get_cache_key_info(key) + + assert result == "a" * 47 + "..." + assert result is not None + assert len(result) == 50 + + +def test_get_cache_key_info_exactly_fifty_chars_is_not_truncated() -> None: + """Should leave a cache key of exactly 50 characters untouched.""" + key = "b" * 50 + assert get_cache_key_info(key) == key diff --git a/tests/unit_tests/mcp_service/utils/test_oauth2_utils.py b/tests/unit_tests/mcp_service/utils/test_oauth2_utils.py new file mode 100644 index 00000000000..b610f877937 --- /dev/null +++ b/tests/unit_tests/mcp_service/utils/test_oauth2_utils.py @@ -0,0 +1,96 @@ +# 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 MCP service OAuth2 utilities. +""" + +import pytest + +from superset.exceptions import OAuth2RedirectError +from superset.mcp_service.utils.oauth2_utils import ( + build_oauth2_redirect_message, + OAUTH2_CONFIG_ERROR_MESSAGE, +) + + +def _make_error( + url: str = "https://example.com/oauth/authorize", + tab_id: str = "tab-123", + redirect_uri: str = "https://example.com/oauth/callback", +) -> OAuth2RedirectError: + return OAuth2RedirectError(url=url, tab_id=tab_id, redirect_uri=redirect_uri) + + +def test_build_oauth2_redirect_message_includes_the_authorization_url() -> None: + """Should embed the OAuth2 authorization URL from the exception's extra.""" + ex = _make_error(url="https://auth.example.com/start") + + message = build_oauth2_redirect_message(ex) + + assert "https://auth.example.com/start" in message + + +def test_build_oauth2_redirect_message_explains_next_steps() -> None: + """Should tell the user to open the URL and retry the request.""" + ex = _make_error() + + message = build_oauth2_redirect_message(ex) + + assert message.startswith("This database uses OAuth for authentication.") + assert "open the following URL in your browser" in message + assert "retry this request" in message + + +def test_build_oauth2_redirect_message_ends_with_url_on_its_own_line() -> None: + """Should place the URL at the end of the message after a blank line.""" + ex = _make_error(url="https://auth.example.com/oauth2/authorize?client_id=abc") + + message = build_oauth2_redirect_message(ex) + + assert message.endswith("https://auth.example.com/oauth2/authorize?client_id=abc") + assert "\n\n" in message + + +def test_build_oauth2_redirect_message_raises_assertion_when_extra_missing() -> None: + """Should assert when the exception's error.extra is unexpectedly None.""" + ex = _make_error() + ex.error.extra = None + + with pytest.raises(AssertionError): + build_oauth2_redirect_message(ex) + + +def test_build_oauth2_redirect_message_raises_key_error_when_url_missing() -> None: + """Should raise a KeyError if extra is present but lacks a 'url' key.""" + ex = _make_error() + ex.error.extra = {"tab_id": "tab-123", "redirect_uri": "https://example.com/cb"} + + with pytest.raises(KeyError): + build_oauth2_redirect_message(ex) + + +def test_oauth2_config_error_message_is_a_nonempty_string() -> None: + """Should expose a non-empty, human-readable config error message.""" + assert isinstance(OAUTH2_CONFIG_ERROR_MESSAGE, str) + assert len(OAUTH2_CONFIG_ERROR_MESSAGE) > 0 + + +def test_oauth2_config_error_message_points_to_the_administrator() -> None: + """Should direct the user to contact their Superset administrator.""" + assert "administrator" in OAUTH2_CONFIG_ERROR_MESSAGE + assert "Superset" in OAUTH2_CONFIG_ERROR_MESSAGE diff --git a/tests/unit_tests/mcp_service/utils/test_retry_utils.py b/tests/unit_tests/mcp_service/utils/test_retry_utils.py new file mode 100644 index 00000000000..9cde4f71410 --- /dev/null +++ b/tests/unit_tests/mcp_service/utils/test_retry_utils.py @@ -0,0 +1,396 @@ +# 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 MCP service retry utilities. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from sqlalchemy.exc import OperationalError + +from superset.mcp_service.utils.retry_utils import ( + async_retry_database_operation, + async_retry_on_exception, + exponential_backoff, + retry_database_operation, + retry_on_exception, + retry_screenshot_operation, + RetryableOperation, +) + +SLEEP: str = "superset.mcp_service.utils.retry_utils.time.sleep" +ASYNC_SLEEP: str = "superset.mcp_service.utils.retry_utils.asyncio.sleep" + + +# --------------------------------------------------------------------------- +# exponential_backoff +# --------------------------------------------------------------------------- + + +def test_exponential_backoff_without_jitter_doubles_each_attempt() -> None: + """Should double the delay for each subsequent attempt when jitter is off.""" + assert exponential_backoff(0, base_delay=1.0, jitter=False) == 1.0 + assert exponential_backoff(1, base_delay=1.0, jitter=False) == 2.0 + assert exponential_backoff(2, base_delay=1.0, jitter=False) == 4.0 + + +def test_exponential_backoff_caps_at_max_delay() -> None: + """Should never return a delay larger than max_delay.""" + delay = exponential_backoff(10, base_delay=1.0, max_delay=5.0, jitter=False) + assert delay == 5.0 + + +def test_exponential_backoff_jitter_stays_within_twenty_five_percent() -> None: + """Should keep jittered delays within +/-25% of the base delay.""" + for _ in range(50): + delay = exponential_backoff(2, base_delay=1.0, jitter=True) + # Unjittered delay for attempt=2, base_delay=1.0 is 4.0. + assert 3.0 <= delay <= 5.0 + + +def test_exponential_backoff_never_negative() -> None: + """Should clamp the delay to be non-negative even with jitter applied.""" + for _ in range(50): + delay = exponential_backoff(0, base_delay=0.001, jitter=True) + assert delay >= 0 + + +# --------------------------------------------------------------------------- +# retry_on_exception (sync) +# --------------------------------------------------------------------------- + + +def test_retry_on_exception_succeeds_first_try() -> None: + """Should call the wrapped function once when it succeeds immediately.""" + mock_func = MagicMock(return_value="ok") + wrapped = retry_on_exception(max_attempts=3)(mock_func) + + result = wrapped() + + assert result == "ok" + assert mock_func.call_count == 1 + + +def test_retry_on_exception_retries_then_succeeds() -> None: + """Should retry on retryable exceptions and return the eventual success.""" + mock_func = MagicMock( + side_effect=[ConnectionError("fail"), ConnectionError("fail"), "ok"] + ) + mock_func.__name__ = "mock_func" + + with patch(SLEEP) as mock_sleep: + wrapped = retry_on_exception(max_attempts=3, base_delay=0.01, jitter=False)( + mock_func + ) + result = wrapped() + + assert result == "ok" + assert mock_func.call_count == 3 + assert mock_sleep.call_count == 2 + + +def test_retry_on_exception_exhausts_retries_and_raises_last_exception() -> None: + """Should raise the last retryable exception once max_attempts is reached.""" + mock_func = MagicMock(side_effect=ConnectionError("always fails")) + mock_func.__name__ = "mock_func" + + with patch(SLEEP) as mock_sleep: + wrapped = retry_on_exception(max_attempts=3, base_delay=0.01, jitter=False)( + mock_func + ) + with pytest.raises(ConnectionError, match="always fails"): + wrapped() + + assert mock_func.call_count == 3 + # No sleep after the final (failed) attempt. + assert mock_sleep.call_count == 2 + + +def test_retry_on_exception_non_retryable_exception_fails_immediately() -> None: + """Should not retry exceptions outside the configured retryable tuple.""" + mock_func = MagicMock(side_effect=ValueError("bad input")) + mock_func.__name__ = "mock_func" + + with patch(SLEEP) as mock_sleep: + wrapped = retry_on_exception(max_attempts=3, exceptions=(ConnectionError,))( + mock_func + ) + with pytest.raises(ValueError, match="bad input"): + wrapped() + + assert mock_func.call_count == 1 + mock_sleep.assert_not_called() + + +def test_retry_on_exception_respects_custom_exceptions_tuple() -> None: + """Should retry only on the exception types passed via `exceptions`.""" + mock_func = MagicMock(side_effect=[TimeoutError("timed out"), "ok"]) + mock_func.__name__ = "mock_func" + + with patch(SLEEP): + wrapped = retry_on_exception( + max_attempts=2, base_delay=0.01, jitter=False, exceptions=(TimeoutError,) + )(mock_func) + result = wrapped() + + assert result == "ok" + assert mock_func.call_count == 2 + + +def test_retry_on_exception_max_attempts_zero_raises_runtime_error() -> None: + """Should raise RuntimeError when max_attempts leaves no attempt to run.""" + mock_func = MagicMock(return_value="ok") + mock_func.__name__ = "mock_func" + wrapped = retry_on_exception(max_attempts=0)(mock_func) + + with pytest.raises(RuntimeError, match="All 0 attempts failed"): + wrapped() + + mock_func.assert_not_called() + + +# --------------------------------------------------------------------------- +# async_retry_on_exception +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_async_retry_on_exception_succeeds_first_try() -> None: + """Should await the wrapped coroutine once when it succeeds immediately.""" + mock_func = AsyncMock(return_value="ok") + wrapped = async_retry_on_exception(max_attempts=3)(mock_func) + + result = await wrapped() + + assert result == "ok" + assert mock_func.await_count == 1 + + +@pytest.mark.asyncio +async def test_async_retry_on_exception_retries_then_succeeds() -> None: + """Should retry a failing coroutine on retryable exceptions until success.""" + mock_func = AsyncMock(side_effect=[ConnectionError("fail"), "ok"]) + + with patch(ASYNC_SLEEP, new_callable=AsyncMock) as mock_sleep: + wrapped = async_retry_on_exception( + max_attempts=3, base_delay=0.01, jitter=False + )(mock_func) + result = await wrapped() + + assert result == "ok" + assert mock_func.await_count == 2 + mock_sleep.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_async_retry_on_exception_exhausts_retries_and_raises() -> None: + """Should raise the last retryable exception once max_attempts is reached.""" + mock_func = AsyncMock(side_effect=ConnectionError("nope")) + + with patch(ASYNC_SLEEP, new_callable=AsyncMock): + wrapped = async_retry_on_exception( + max_attempts=2, base_delay=0.01, jitter=False + )(mock_func) + with pytest.raises(ConnectionError, match="nope"): + await wrapped() + + assert mock_func.await_count == 2 + + +@pytest.mark.asyncio +async def test_async_retry_on_exception_non_retryable_fails_immediately() -> None: + """Should not retry exceptions outside the configured retryable tuple.""" + mock_func = AsyncMock(side_effect=ValueError("bad")) + + with patch(ASYNC_SLEEP, new_callable=AsyncMock) as mock_sleep: + wrapped = async_retry_on_exception( + max_attempts=3, exceptions=(ConnectionError,) + )(mock_func) + with pytest.raises(ValueError, match="bad"): + await wrapped() + + assert mock_func.await_count == 1 + mock_sleep.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# RetryableOperation +# --------------------------------------------------------------------------- + + +def test_retryable_operation_success_does_not_suppress_or_retry() -> None: + """Should leave attempt count untouched when the block succeeds.""" + op = RetryableOperation("test_op", max_attempts=3) + + with op: + pass + + assert op.current_attempt == 0 + assert op.last_exception is None + + +def test_retryable_operation_retries_until_success() -> None: + """Should suppress retryable exceptions and allow the loop to retry.""" + op = RetryableOperation("test_op", max_attempts=3, base_delay=0.01, jitter=False) + attempts: int = 0 + result: str | None = None + + with patch(SLEEP) as mock_sleep: + while op.should_retry(): + with op: + attempts += 1 + if attempts < 3: + raise ConnectionError("fail") + result = "success" + break + + assert result == "success" + assert attempts == 3 + assert mock_sleep.call_count == 2 + + +def _run_always_failing_operation(op: RetryableOperation) -> None: + while op.should_retry(): + with op: + raise ConnectionError("boom") + + +def test_retryable_operation_propagates_after_max_attempts() -> None: + """Should stop suppressing once max_attempts retryable failures occur.""" + op = RetryableOperation("test_op", max_attempts=2, base_delay=0.01, jitter=False) + + with patch(SLEEP), pytest.raises(ConnectionError, match="boom"): + _run_always_failing_operation(op) + + assert op.current_attempt == 2 + assert isinstance(op.last_exception, ConnectionError) + + +def test_retryable_operation_non_retryable_exception_propagates_immediately() -> None: + """Should not suppress exceptions outside the configured retryable tuple.""" + op = RetryableOperation("test_op", max_attempts=3) + + with pytest.raises(ValueError, match="bad"): + with op: + raise ValueError("bad") + + # Non-retryable exceptions don't count as a tracked attempt. + assert op.current_attempt == 0 + + +def test_retryable_operation_should_retry_reflects_attempt_count() -> None: + """Should report should_retry() based on current_attempt vs max_attempts.""" + op = RetryableOperation("test_op", max_attempts=3) + assert op.should_retry() is True + + op.current_attempt = 3 + assert op.should_retry() is False + + +# --------------------------------------------------------------------------- +# Convenience functions +# --------------------------------------------------------------------------- + + +def test_retry_database_operation_success() -> None: + """Should pass through args/kwargs and return the function's result.""" + func = MagicMock(return_value=42) + + result = retry_database_operation(func, "a", max_attempts=2, b=1) + + assert result == 42 + func.assert_called_once_with("a", b=1) + + +def test_retry_database_operation_retries_on_operational_error() -> None: + """Should retry on sqlalchemy OperationalError and return eventual success.""" + func = MagicMock( + side_effect=[OperationalError("stmt", {}, Exception("orig")), "ok"] + ) + + with patch(SLEEP): + result = retry_database_operation(func, max_attempts=2) + + assert result == "ok" + assert func.call_count == 2 + + +def test_retry_database_operation_does_not_retry_non_retryable_exceptions() -> None: + """Should fail immediately for exceptions outside the retryable set.""" + func = MagicMock(side_effect=ValueError("bad")) + + with pytest.raises(ValueError, match="bad"): + retry_database_operation(func, max_attempts=3) + + assert func.call_count == 1 + + +@pytest.mark.asyncio +async def test_async_retry_database_operation_success() -> None: + """Should await the coroutine and return its result.""" + func = AsyncMock(return_value="ok") + + result = await async_retry_database_operation(func, max_attempts=2) + + assert result == "ok" + func.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_async_retry_database_operation_retries_on_timeout_error() -> None: + """Should retry an async database operation on sqlalchemy TimeoutError.""" + from sqlalchemy.exc import TimeoutError as SATimeoutError + + func = AsyncMock(side_effect=[SATimeoutError("timed out"), "ok"]) + + with patch(ASYNC_SLEEP, new_callable=AsyncMock): + result = await async_retry_database_operation(func, max_attempts=2) + + assert result == "ok" + assert func.await_count == 2 + + +def test_retry_screenshot_operation_success() -> None: + """Should return the screenshot function's result on success.""" + func = MagicMock(return_value=b"png-bytes") + + result = retry_screenshot_operation(func) + + assert result == b"png-bytes" + + +def test_retry_screenshot_operation_retries_on_os_error() -> None: + """Should retry a screenshot operation on OSError.""" + func = MagicMock(side_effect=[OSError("disk full"), b"png-bytes"]) + + with patch(SLEEP): + result = retry_screenshot_operation(func, max_attempts=2) + + assert result == b"png-bytes" + assert func.call_count == 2 + + +def test_retry_screenshot_operation_does_not_retry_non_retryable_exceptions() -> None: + """Should fail immediately for exceptions outside the screenshot retry set.""" + func = MagicMock(side_effect=ValueError("bad")) + + with pytest.raises(ValueError, match="bad"): + retry_screenshot_operation(func, max_attempts=2) + + assert func.call_count == 1