tests/unit_tests/thumbnails/test_digest.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.
from __future__ import annotations

from contextlib import nullcontext
from typing import Any, TYPE_CHECKING
from unittest.mock import MagicMock, patch, PropertyMock

import pytest
from flask_appbuilder.security.sqla.models import User

from superset.connectors.sqla.models import BaseDatasource, SqlaTable
from superset.tasks.exceptions import InvalidExecutorError
from superset.tasks.types import Executor, ExecutorType, FixedExecutor
from superset.utils.core import DatasourceType, override_user

if TYPE_CHECKING:
    from superset.models.dashboard import Dashboard
    from superset.models.slice import Slice

_DEFAULT_DASHBOARD_KWARGS: dict[str, Any] = {
    "id": 1,
    "dashboard_title": "My Title",
    "slices": [{"id": 1, "slice_name": "My Chart"}],
    "position_json": '{"a": "b"}',
    "css": "background-color: lightblue;",
    "json_metadata": '{"c": "d"}',
}

_DEFAULT_CHART_KWARGS = {
    "id": 2,
    "params": {"a": "b"},
}


def CUSTOM_DASHBOARD_FUNC(  # noqa: N802
    dashboard: Dashboard,
    executor_type: ExecutorType,
    executor: str,
) -> str:
    return f"{dashboard.id}.{executor_type.value}.{executor}"


def CUSTOM_CHART_FUNC(  # noqa: N802
    chart: Slice,
    executor_type: ExecutorType,
    executor: str,
) -> str:
    return f"{chart.id}.{executor_type.value}.{executor}"


def prepare_datasource_mock(
    datasource_conf: dict[str, Any], spec: type[BaseDatasource | SqlaTable]
) -> BaseDatasource | SqlaTable:
    datasource = MagicMock(spec=spec)
    datasource.id = 1
    datasource.type = DatasourceType.TABLE
    datasource.is_rls_supported = datasource_conf.get("is_rls_supported", False)
    datasource.get_sqla_row_level_filters = datasource_conf.get(
        "get_sqla_row_level_filters", MagicMock(return_value=[])
    )
    return datasource


@pytest.mark.parametrize(
    "dashboard_overrides,execute_as,has_current_user,use_custom_digest,rls_datasources,expected_result",
    [
        (
            None,
            [FixedExecutor("admin")],
            False,
            False,
            [],
            "71452fee8ffbd8d340193d611bcd4559",
        ),
        (
            None,
            [ExecutorType.CURRENT_USER],
            True,
            False,
            [],
            "209dc060ac19271b8708731e3b8280f5",
        ),
        (
            {
                "dashboard_title": "My Other Title",
            },
            [ExecutorType.CURRENT_USER],
            True,
            False,
            [],
            "209dc060ac19271b8708731e3b8280f5",
        ),
        (
            {
                "id": 2,
            },
            [ExecutorType.CURRENT_USER],
            True,
            False,
            [],
            "06a4144466dbd5ffad0c3c2225e96296",
        ),
        (
            {
                "slices": [{"id": 2, "slice_name": "My Other Chart"}],
            },
            [ExecutorType.CURRENT_USER],
            True,
            False,
            [],
            "a823ece9563895ccb14f3d9095e84f7a",
        ),
        (
            {
                "position_json": {"b": "c"},
            },
            [ExecutorType.CURRENT_USER],
            True,
            False,
            [],
            "33c5475f92a904925ab3ef493526e5b5",
        ),
        (
            {
                "css": "background-color: darkblue;",
            },
            [ExecutorType.CURRENT_USER],
            True,
            False,
            [],
            "cec57345e6402c0d4b3caee5cfaa0a03",
        ),
        (
            {
                "json_metadata": {"d": "e"},
            },
            [ExecutorType.CURRENT_USER],
            True,
            False,
            [],
            "5380dcbe94621a0759b09554404f3d02",
        ),
        (
            None,
            [ExecutorType.CURRENT_USER],
            True,
            False,
            [
                {
                    "is_rls_supported": True,
                    "get_sqla_row_level_filters": MagicMock(return_value=["filter1"]),
                }
            ],
            "4138959f275c1991466cafcfb190fd72",
        ),
        (
            None,
            [ExecutorType.CURRENT_USER],
            True,
            False,
            [
                {
                    "is_rls_supported": True,
                    "get_sqla_row_level_filters": MagicMock(
                        return_value=["filter1", "filter2"]
                    ),
                },
                {
                    "is_rls_supported": True,
                    "get_sqla_row_level_filters": MagicMock(
                        return_value=["filter3", "filter4"]
                    ),
                },
            ],
            "80d3bfcc7144bccdba8c718cf49b6420",
        ),
        (
            None,
            [ExecutorType.CURRENT_USER],
            True,
            False,
            [
                {
                    "is_rls_supported": False,
                    "get_sqla_row_level_filters": MagicMock(return_value=[]),
                },
                {
                    "is_rls_supported": True,
                    "get_sqla_row_level_filters": MagicMock(
                        return_value=["filter1", "filter2"]
                    ),
                },
            ],
            "e8fc68cd5aba22a5f1acf06164bfc0f4",
        ),
        (
            None,
            [ExecutorType.CURRENT_USER],
            False,
            False,
            [],
            None,
        ),
        (
            None,
            [ExecutorType.FIXED_USER],
            False,
            False,
            [],
            InvalidExecutorError(),
        ),
    ],
)
def test_dashboard_digest(
    dashboard_overrides: dict[str, Any] | None,
    execute_as: list[Executor],
    has_current_user: bool,
    use_custom_digest: bool,
    rls_datasources: list[dict[str, Any]],
    expected_result: str | Exception,
) -> None:
    from superset import app, security_manager
    from superset.models.dashboard import Dashboard
    from superset.models.slice import Slice
    from superset.thumbnails.digest import get_dashboard_digest

    # Prepare dashboard and slices
    kwargs = {
        **_DEFAULT_DASHBOARD_KWARGS,
        **(dashboard_overrides or {}),
    }
    slices = [Slice(**slice_kwargs) for slice_kwargs in kwargs.pop("slices")]
    dashboard = Dashboard(**kwargs, slices=slices)

    # Mock datasources with RLS
    datasources = []
    for rls_source in rls_datasources:
        datasource = prepare_datasource_mock(rls_source, BaseDatasource)
        datasources.append(datasource)

    user: User | None = None
    if has_current_user:
        user = User(id=1, username="1")

    func = CUSTOM_DASHBOARD_FUNC if use_custom_digest else None

    with (
        patch.dict(
            app.config,
            {
                "THUMBNAIL_EXECUTORS": execute_as,
                "THUMBNAIL_DASHBOARD_DIGEST_FUNC": func,
            },
        ),
        patch.object(
            type(dashboard),
            "datasources",
            new_callable=PropertyMock,
            return_value=datasources,
        ),
        patch.object(security_manager, "find_user", return_value=user),
        override_user(user),
    ):
        cm = (
            pytest.raises(type(expected_result))
            if isinstance(expected_result, Exception)
            else nullcontext()
        )
        with cm:
            assert get_dashboard_digest(dashboard=dashboard) == expected_result


@pytest.mark.parametrize(
    "chart_overrides,execute_as,has_current_user,use_custom_digest,rls_datasource,expected_result",
    [
        (
            None,
            [FixedExecutor("admin")],
            False,
            False,
            None,
            "47d852b5c4df211c115905617bb722c1",
        ),
        (
            None,
            [ExecutorType.CURRENT_USER],
            True,
            False,
            None,
            "4f8109d3761e766e650af514bb358f10",
        ),
        (
            None,
            [ExecutorType.CURRENT_USER],
            True,
            True,
            None,
            "2.current_user.1",
        ),
        (
            None,
            [ExecutorType.CURRENT_USER],
            True,
            False,
            {
                "is_rls_supported": True,
                "get_sqla_row_level_filters": MagicMock(return_value=["filter1"]),
            },
            "61e70336c27eb97fb050328a0b050373",
        ),
        (
            None,
            [ExecutorType.CURRENT_USER],
            True,
            False,
            {
                "is_rls_supported": True,
                "get_sqla_row_level_filters": MagicMock(
                    return_value=["filter1", "filter2"]
                ),
            },
            "95c7cefde8cb519f005f33bfb33cb196",
        ),
        (
            None,
            [ExecutorType.CURRENT_USER],
            True,
            False,
            {
                "is_rls_supported": False,
                "get_sqla_row_level_filters": MagicMock(return_value=[]),
            },
            "4f8109d3761e766e650af514bb358f10",
        ),
        (
            None,
            [ExecutorType.CURRENT_USER],
            False,
            False,
            None,
            None,
        ),
        (
            None,
            [ExecutorType.FIXED_USER],
            False,
            False,
            None,
            InvalidExecutorError(),
        ),
    ],
)
def test_chart_digest(
    chart_overrides: dict[str, Any] | None,
    execute_as: list[Executor],
    has_current_user: bool,
    use_custom_digest: bool,
    rls_datasource: dict[str, Any] | None,
    expected_result: str | Exception,
) -> None:
    from superset import app, security_manager
    from superset.models.slice import Slice
    from superset.thumbnails.digest import get_chart_digest

    # Mock datasource with RLS if provided
    datasource = None
    if rls_datasource:
        datasource = prepare_datasource_mock(rls_datasource, SqlaTable)

    # Prepare chart with the datasource in the constructor
    kwargs = {
        **_DEFAULT_CHART_KWARGS,
        **(chart_overrides or {}),
    }
    chart = Slice(**kwargs)

    user: User | None = None
    if has_current_user:
        user = User(id=1, username="1")

    func = CUSTOM_CHART_FUNC if use_custom_digest else None

    with (
        patch.dict(
            app.config,
            {
                "THUMBNAIL_EXECUTORS": execute_as,
                "THUMBNAIL_CHART_DIGEST_FUNC": func,
            },
        ),
        patch.object(
            type(chart),
            "datasource",
            new_callable=PropertyMock,
            return_value=datasource,
        ),
        patch.object(security_manager, "find_user", return_value=user),
        override_user(user),
    ):
        cm = (
            pytest.raises(type(expected_result))
            if isinstance(expected_result, Exception)
            else nullcontext()
        )
        with cm:
            assert get_chart_digest(chart=chart) == expected_result