superset2/tests/unit_tests/mcp_service/test_mcp_storage.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 storage factory."""

from unittest.mock import MagicMock, patch


def test_get_mcp_store_returns_none_when_disabled():
    """Storage returns None when MCP_STORE_CONFIG.enabled is False."""
    mock_flask_app = MagicMock()
    mock_flask_app.config.get.return_value = {"enabled": False}

    with patch(
        "superset.mcp_service.flask_singleton.get_flask_app",
        return_value=mock_flask_app,
    ):
        with patch("flask.has_app_context", return_value=True):
            from superset.mcp_service.storage import get_mcp_store

            result = get_mcp_store(prefix="test_")

            assert result is None


def test_get_mcp_store_returns_none_when_no_redis_url():
    """Storage returns None when CACHE_REDIS_URL is not configured."""
    mock_flask_app = MagicMock()
    mock_flask_app.config.get.return_value = {
        "enabled": True,
        "CACHE_REDIS_URL": None,
        "WRAPPER_TYPE": "key_value.aio.wrappers.prefix_keys.PrefixKeysWrapper",
    }

    with patch(
        "superset.mcp_service.flask_singleton.get_flask_app",
        return_value=mock_flask_app,
    ):
        with patch("flask.has_app_context", return_value=True):
            from superset.mcp_service.storage import get_mcp_store

            result = get_mcp_store(prefix="test_")

            assert result is None


def test_get_mcp_store_creates_store_when_enabled():
    """Storage creates wrapped RedisStore when properly configured."""
    mock_flask_app = MagicMock()
    mock_flask_app.config.get.return_value = {
        "enabled": True,
        "CACHE_REDIS_URL": "redis://localhost:6379/0",
        "WRAPPER_TYPE": "key_value.aio.wrappers.prefix_keys.PrefixKeysWrapper",
    }

    mock_redis_store = MagicMock()
    mock_redis_client = MagicMock()
    mock_wrapper_instance = MagicMock()
    mock_wrapper_class = MagicMock(return_value=mock_wrapper_instance)

    with patch(
        "superset.mcp_service.flask_singleton.get_flask_app",
        return_value=mock_flask_app,
    ):
        with patch("flask.has_app_context", return_value=True):
            with patch(
                "superset.mcp_service.storage._import_wrapper_class",
                return_value=mock_wrapper_class,
            ):
                with patch(
                    "key_value.aio.stores.redis.RedisStore",
                    return_value=mock_redis_store,
                ):
                    with patch(
                        "superset.mcp_service.storage.Redis",
                        return_value=mock_redis_client,
                    ):
                        from superset.mcp_service.storage import get_mcp_store

                        result = get_mcp_store(prefix="test_prefix_")

                        # Verify store was created
                        assert result is mock_wrapper_instance
                        # Verify wrapper was called with correct args
                        mock_wrapper_class.assert_called_once_with(
                            key_value=mock_redis_store, prefix="test_prefix_"
                        )


def test_create_redis_store_wrap_false_returns_raw_store():
    """_create_redis_store with wrap=False returns unwrapped RedisStore."""
    store_config = {
        "CACHE_REDIS_URL": "redis://localhost:6379/0",
        "WRAPPER_TYPE": "key_value.aio.wrappers.prefix_keys.PrefixKeysWrapper",
    }

    mock_redis_store = MagicMock()
    mock_redis_client = MagicMock()

    with patch(
        "key_value.aio.stores.redis.RedisStore",
        return_value=mock_redis_store,
    ) as mock_redis_store_class:
        with patch(
            "superset.mcp_service.storage.Redis",
            return_value=mock_redis_client,
        ) as mock_redis_class:
            from superset.mcp_service.storage import _create_redis_store

            result = _create_redis_store(store_config, wrap=False)

            # Verify raw store is returned (not wrapped)
            assert result is mock_redis_store
            # Verify Redis client was created with correct params
            mock_redis_class.assert_called_once()
            call_kwargs = mock_redis_class.call_args[1]
            assert call_kwargs["host"] == "localhost"
            assert call_kwargs["port"] == 6379
            # Verify RedisStore was called with the client
            mock_redis_store_class.assert_called_once_with(client=mock_redis_client)


def test_create_redis_store_wrap_true_requires_prefix():
    """_create_redis_store with wrap=True requires prefix parameter."""
    store_config = {
        "CACHE_REDIS_URL": "redis://localhost:6379/0",
        "WRAPPER_TYPE": "key_value.aio.wrappers.prefix_keys.PrefixKeysWrapper",
    }

    mock_redis_store = MagicMock()

    with patch(
        "key_value.aio.stores.redis.RedisStore",
        return_value=mock_redis_store,
    ):
        from superset.mcp_service.storage import _create_redis_store

        # wrap=True (default) with no prefix should return None
        result = _create_redis_store(store_config, prefix=None, wrap=True)

        assert result is None


def test_create_redis_store_handles_ssl_url():
    """_create_redis_store handles rediss:// URLs with SSL configuration."""
    store_config = {
        "CACHE_REDIS_URL": "rediss://:password@redis.example.com:6380/1",
        "WRAPPER_TYPE": "key_value.aio.wrappers.prefix_keys.PrefixKeysWrapper",
    }

    mock_redis_store = MagicMock()
    mock_redis_client = MagicMock()

    with patch(
        "key_value.aio.stores.redis.RedisStore",
        return_value=mock_redis_store,
    ):
        with patch(
            "superset.mcp_service.storage.Redis",
            return_value=mock_redis_client,
        ) as mock_redis_class:
            from superset.mcp_service.storage import _create_redis_store

            result = _create_redis_store(store_config, wrap=False)

            # Verify store was created
            assert result is mock_redis_store
            # Verify Redis client was created with SSL params
            call_kwargs = mock_redis_class.call_args[1]
            assert call_kwargs["ssl"] is True
            assert call_kwargs["ssl_cert_reqs"] == "none"
            assert call_kwargs["host"] == "redis.example.com"
            assert call_kwargs["port"] == 6380
            assert call_kwargs["db"] == 1


def test_create_redis_store_non_ssl_url_no_ssl_param():
    """_create_redis_store with redis:// URL doesn't pass SSL params."""
    store_config = {
        "CACHE_REDIS_URL": "redis://localhost:6379/0",
    }

    mock_redis_store = MagicMock()
    mock_redis_client = MagicMock()

    with patch(
        "key_value.aio.stores.redis.RedisStore",
        return_value=mock_redis_store,
    ):
        with patch(
            "superset.mcp_service.storage.Redis",
            return_value=mock_redis_client,
        ) as mock_redis_class:
            from superset.mcp_service.storage import _create_redis_store

            result = _create_redis_store(store_config, wrap=False)

            assert result is mock_redis_store
            # Verify SSL params were NOT passed for non-SSL URL
            call_kwargs = mock_redis_class.call_args[1]
            assert "ssl" not in call_kwargs
            assert "ssl_cert_reqs" not in call_kwargs


def test_create_redis_store_handles_url_with_username_and_password():
    """_create_redis_store properly handles URL with username and password."""
    test_password = "mypassword"  # noqa: S105
    store_config = {
        "CACHE_REDIS_URL": f"redis://myuser:{test_password}@redis.example.com:6379/0",
    }

    mock_redis_store = MagicMock()
    mock_redis_client = MagicMock()

    with patch(
        "key_value.aio.stores.redis.RedisStore",
        return_value=mock_redis_store,
    ):
        with patch(
            "superset.mcp_service.storage.Redis",
            return_value=mock_redis_client,
        ) as mock_redis_class:
            from superset.mcp_service.storage import _create_redis_store

            result = _create_redis_store(store_config, wrap=False)

            assert result is mock_redis_store
            # Verify Redis client was created with password from URL
            call_kwargs = mock_redis_class.call_args[1]
            assert call_kwargs["host"] == "redis.example.com"
            assert call_kwargs["password"] == test_password


def test_create_redis_store_handles_url_with_only_username():
    """_create_redis_store handles URL with username but no password."""
    store_config = {
        "CACHE_REDIS_URL": "redis://myuser@redis.example.com:6379/0",
    }

    mock_redis_store = MagicMock()
    mock_redis_client = MagicMock()

    with patch(
        "key_value.aio.stores.redis.RedisStore",
        return_value=mock_redis_store,
    ):
        with patch(
            "superset.mcp_service.storage.Redis",
            return_value=mock_redis_client,
        ) as mock_redis_class:
            from superset.mcp_service.storage import _create_redis_store

            result = _create_redis_store(store_config, wrap=False)

            assert result is mock_redis_store
            # Verify Redis client was created with correct params
            call_kwargs = mock_redis_class.call_args[1]
            assert call_kwargs["host"] == "redis.example.com"
            # No password in URL means password should be None
            assert call_kwargs["password"] is None