superset2/tests/unit_tests/mcp_service/test_mcp_config.py

# 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.

"""Tests for MCP service configuration and branding."""

from unittest.mock import MagicMock, patch

from superset.mcp_service.app import get_default_instructions, init_fastmcp_server


def test_get_default_instructions_with_default_branding():
    """Test that default branding produces Apache Superset in instructions."""
    instructions = get_default_instructions()

    assert "Apache Superset" in instructions
    assert "Apache Superset MCP" in instructions
    assert "model context protocol" in instructions.lower()


def test_get_default_instructions_with_custom_branding():
    """Test that custom branding is reflected in instructions."""
    custom_branding = "ACME Analytics"
    instructions = get_default_instructions(branding=custom_branding)

    assert custom_branding in instructions
    assert f"{custom_branding} MCP" in instructions
    # Should not contain default Apache Superset branding
    assert "Apache Superset" not in instructions


def test_get_default_instructions_with_enterprise_branding():
    """Test instructions with enterprise/white-label branding."""
    enterprise_branding = "DataViz Platform"
    instructions = get_default_instructions(branding=enterprise_branding)

    assert enterprise_branding in instructions
    assert f"{enterprise_branding} MCP" in instructions
    # Verify it contains expected tool documentation
    assert "list_dashboards" in instructions
    assert "list_charts" in instructions
    assert "execute_sql" in instructions


def test_get_default_instructions_mentions_feature_availability():
    """Test that instructions direct LLMs to get_instance_info for features."""
    instructions = get_default_instructions()

    assert "get_instance_info" in instructions
    assert "Feature Availability" in instructions
    assert "accessible menus" in instructions


def test_get_default_instructions_forbid_disclosing_other_user_access_or_roles() -> (
    None
):
    """Test that instructions route access-list questions to workspace admins."""
    instructions = get_default_instructions()

    assert "Do NOT disclose dashboard access lists" in instructions
    assert "other users' names, usernames, email addresses" in instructions
    assert "current user's own identity details" in instructions
    assert "Do NOT use execute_sql to query user, role, owner" in instructions
    assert "direct them to their workspace admin" in instructions


def test_init_fastmcp_server_with_default_app_name():
    """Test that default APP_NAME produces Superset branding."""
    # Mock Flask app config with default APP_NAME
    mock_flask_app = MagicMock()
    mock_flask_app.config.get.return_value = "Superset"

    # Patch at the import location to avoid actual Flask app creation
    with patch.dict(
        "sys.modules",
        {"superset.mcp_service.flask_singleton": MagicMock(app=mock_flask_app)},
    ):
        with patch("superset.mcp_service.app.mcp") as mock_mcp:
            init_fastmcp_server()

            # Verify the global mcp instance was configured with Superset branding
            assert "Superset MCP" in mock_mcp._mcp_server.instructions
            assert "Superset dashboards" in mock_mcp._mcp_server.instructions


def test_init_fastmcp_server_with_custom_app_name():
    """Test that custom APP_NAME produces branded instructions."""
    custom_app_name = "ACME Analytics"
    # Mock Flask app config with custom APP_NAME
    mock_flask_app = MagicMock()
    mock_flask_app.config.get.return_value = custom_app_name

    # Patch at the import location to avoid actual Flask app creation
    with patch.dict(
        "sys.modules",
        {"superset.mcp_service.flask_singleton": MagicMock(app=mock_flask_app)},
    ):
        with patch("superset.mcp_service.app.mcp") as mock_mcp:
            init_fastmcp_server()

            # Verify instructions use custom branding
            assert custom_app_name in mock_mcp._mcp_server.instructions
            # Should not contain default Apache Superset branding
            assert "Apache Superset" not in mock_mcp._mcp_server.instructions


def test_init_fastmcp_server_derives_server_name_from_app_name():
    """Test that server name is derived from APP_NAME."""
    custom_app_name = "DataViz Platform"
    expected_server_name = f"{custom_app_name} MCP Server"

    # Mock Flask app config
    mock_flask_app = MagicMock()
    mock_flask_app.config.get.return_value = custom_app_name

    # Patch at the import location to avoid actual Flask app creation
    with patch.dict(
        "sys.modules",
        {"superset.mcp_service.flask_singleton": MagicMock(app=mock_flask_app)},
    ):
        with patch("superset.mcp_service.app.mcp") as mock_mcp:
            init_fastmcp_server()

            # Verify the global mcp instance got the derived name
            assert mock_mcp._mcp_server.name == expected_server_name


def test_init_fastmcp_server_applies_auth_to_global_instance():
    """Test that auth is applied to the global mcp instance, not a new one."""
    mock_flask_app = MagicMock()
    mock_flask_app.config.get.return_value = "Superset"
    mock_auth = MagicMock()

    with patch.dict(
        "sys.modules",
        {"superset.mcp_service.flask_singleton": MagicMock(app=mock_flask_app)},
    ):
        with patch("superset.mcp_service.app.mcp") as mock_mcp:
            result = init_fastmcp_server(auth=mock_auth)

            # Auth should be set on the global instance
            assert mock_mcp.auth == mock_auth
            # Should return the global instance (not a new one)
            assert result is mock_mcp


def test_init_fastmcp_server_applies_middleware_to_global_instance():
    """Test that middleware is added to the global mcp instance."""
    mock_flask_app = MagicMock()
    mock_flask_app.config.get.return_value = "Superset"
    mock_mw = MagicMock()

    with patch.dict(
        "sys.modules",
        {"superset.mcp_service.flask_singleton": MagicMock(app=mock_flask_app)},
    ):
        with patch("superset.mcp_service.app.mcp") as mock_mcp:
            init_fastmcp_server(middleware=[mock_mw])

            # Middleware should be added via add_middleware
            mock_mcp.add_middleware.assert_called_once_with(mock_mw)