superset2/tests/unit_tests/subjects/test_utils.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 unittest.mock import MagicMock, patch

import pytest

from superset.commands.utils import (
    compute_subject_list,
    compute_subjects,
    populate_subject_list,
    populate_subjects,
)
from superset.subjects.exceptions import SubjectsNotFoundValidationError
from superset.subjects.models import Subject
from superset.subjects.types import SubjectType
from superset.subjects.utils import (
    get_current_user_subject_ids,
    get_user_subject_ids_subquery,
)


def _make_subject(
    id_: int,
    type_: SubjectType = SubjectType.USER,
    user: MagicMock | None = None,
) -> Subject:
    s = Subject()
    s.id = id_
    s.type = type_
    s.user = user
    return s


def _mock_sm(*, is_admin: bool = False) -> MagicMock:
    sm = MagicMock()
    sm.is_admin = MagicMock(return_value=is_admin)
    return sm


def test_get_user_subject_ids_subquery_includes_group_roles(app_context):
    """User subjects include roles inherited through group membership."""
    query = get_user_subject_ids_subquery(7)
    sql = str(query.compile(compile_kwargs={"literal_binds": True}))

    assert "ab_group_role" in sql
    assert "ab_user_group" in sql


def test_get_current_user_subject_ids_guest_user(app_context):
    """Guest users resolve subject IDs from their embedded roles."""
    from flask import g

    role_subjects = [
        _make_subject(10, SubjectType.ROLE),
        _make_subject(20, SubjectType.ROLE),
    ]
    guest_user = MagicMock()
    guest_user.is_guest_user = True
    guest_user.roles = [1, 2]
    g.user = guest_user

    with (
        patch(
            "superset.subjects.utils.subjects_from_roles",
            return_value=role_subjects,
        ) as mock_subjects_from_roles,
        patch("superset.utils.core.get_user_id") as mock_get_user_id,
    ):
        assert get_current_user_subject_ids() == [10, 20]

    mock_subjects_from_roles.assert_called_once_with([1, 2])
    mock_get_user_id.assert_not_called()


# --------------------------------------------------------------------------
# populate_subject_list
# --------------------------------------------------------------------------


@pytest.mark.parametrize(
    "subject_ids,default_to_user,user_subject_exists,expected_count",
    [
        ([], True, True, 1),  # empty + default → user's subject
        (None, True, True, 1),  # None + default → user's subject
        ([], True, False, 0),  # empty + default but no subject → empty
    ],
)
@patch("superset.commands.utils.get_user_subject")
@patch("superset.commands.utils.get_user_id")
def test_populate_subject_list_defaults(
    mock_user_id,
    mock_get_user_subject,
    subject_ids,
    default_to_user,
    user_subject_exists,
    expected_count,
):
    """Default-to-user behavior with various combinations of input."""
    mock_user_id.return_value = 1
    mock_get_user_subject.return_value = (
        _make_subject(10) if user_subject_exists else None
    )

    result = populate_subject_list(subject_ids, default_to_user=default_to_user)
    assert len(result) == expected_count


@patch("superset.commands.utils.get_subject")
@patch("superset.commands.utils.security_manager", new_callable=lambda: _mock_sm)
@patch("superset.commands.utils.get_user_id")
def test_populate_subject_list_admin_vs_non_admin(
    mock_user_id,
    mock_sm,
    mock_get_subject,
):
    """Admin can set arbitrary subjects; non-admin without lockout behaves
    the same when they're not removing themselves."""
    subject_a = _make_subject(20)
    mock_user_id.return_value = 1
    mock_get_subject.return_value = subject_a

    # Admin: only the requested subject
    mock_sm.is_admin = MagicMock(return_value=True)
    result = populate_subject_list([20], default_to_user=False)
    assert result == [subject_a]

    # Non-admin without lockout: same result (no prepend)
    mock_sm.is_admin = MagicMock(return_value=False)
    result = populate_subject_list([20], default_to_user=False, ensure_no_lockout=False)
    assert result == [subject_a]

    # Admin with empty list
    mock_sm.is_admin = MagicMock(return_value=True)
    result = populate_subject_list([], default_to_user=False)
    assert result == []


@patch("superset.commands.utils.get_subject")
@patch("superset.commands.utils.get_user_subject_ids")
@patch("superset.commands.utils.get_user_subject")
@patch("superset.commands.utils.security_manager", new_callable=lambda: _mock_sm)
@patch("superset.commands.utils.get_user_id")
def test_populate_subject_list_lockout_prevention(
    mock_user_id,
    mock_sm,
    mock_get_user_subject,
    mock_get_user_subject_ids,
    mock_get_subject,
):
    """Non-admins cannot remove themselves when ensure_no_lockout is True."""
    user_subject = _make_subject(10)
    other_subject = _make_subject(20)
    mock_user_id.return_value = 1
    mock_sm.is_admin = MagicMock(return_value=False)
    mock_get_user_subject.return_value = user_subject
    mock_get_user_subject_ids.return_value = [10]
    mock_get_subject.return_value = other_subject

    result = populate_subject_list(
        [20],
        default_to_user=False,
        ensure_no_lockout=True,
    )
    # user_subject prepended because 10 not in [20]
    assert len(result) == 2
    assert result[0] == user_subject
    assert result[1] == other_subject


def test_populate_subject_list_extra_editors_resolver_skips_lockout_prevention():
    user_subject = _make_subject(10)
    other_subject = _make_subject(20)
    mock_sm = _mock_sm(is_admin=False)

    with (
        patch("superset.commands.utils.get_user_id", return_value=1),
        patch("superset.commands.utils.security_manager", mock_sm),
        patch("superset.commands.utils.get_user_subject", return_value=user_subject),
        patch("superset.commands.utils.get_user_subject_ids", return_value=[10]),
        patch("superset.commands.utils.get_subject", return_value=other_subject),
        patch(
            "superset.commands.utils._has_extra_editors_resolver",
            return_value=True,
        ),
    ):
        result = populate_subject_list(
            [20],
            default_to_user=False,
            ensure_no_lockout=True,
        )

    assert result == [other_subject]


def test_populate_subject_list_extra_editors_resolver_keeps_default_to_user():
    user_subject = _make_subject(10)

    with (
        patch("superset.commands.utils.get_user_id", return_value=1),
        patch("superset.commands.utils.get_user_subject", return_value=user_subject),
        patch(
            "superset.commands.utils._has_extra_editors_resolver",
            return_value=True,
        ),
    ):
        result = populate_subject_list([], default_to_user=True)

    assert result == [user_subject]


@patch("superset.commands.utils.get_subject")
@patch("superset.commands.utils.security_manager", new_callable=lambda: _mock_sm)
@patch("superset.commands.utils.get_user_id")
def test_populate_subject_list_not_found(mock_user_id, mock_sm, mock_get_subject):
    """Raises SubjectsNotFoundValidationError when a subject ID doesn't exist."""
    mock_user_id.return_value = 1
    mock_sm.is_admin = MagicMock(return_value=True)
    mock_get_subject.return_value = None

    with pytest.raises(SubjectsNotFoundValidationError):
        populate_subject_list([999], default_to_user=False)


# --------------------------------------------------------------------------
# compute_subject_list
# --------------------------------------------------------------------------


@pytest.mark.parametrize(
    "current_ids,new_ids,expected_call_ids",
    [
        ([1, 2], [3, 4], [3, 4]),  # new replaces current
        ([1, 2], None, [1, 2]),  # None preserves current
        ([1, 2], [], []),  # empty clears
        (None, [3, 4], [3, 4]),  # None current + new ids
    ],
    ids=["replace", "preserve", "clear", "none-current"],
)
@patch("superset.commands.utils.populate_subject_list")
def test_compute_subject_list(mock_populate, current_ids, new_ids, expected_call_ids):
    """Delegates to populate_subject_list with the right IDs."""
    current = [_make_subject(i) for i in current_ids] if current_ids else current_ids
    compute_subject_list(current, new_ids)
    mock_populate.assert_called_once_with(
        expected_call_ids,
        default_to_user=False,
        ensure_no_lockout=False,
        field_name="subjects",
    )


# --------------------------------------------------------------------------
# populate_subjects / compute_subjects (consolidated helpers)
# --------------------------------------------------------------------------


@patch("superset.commands.utils.populate_subject_list")
def test_populate_subjects_all_variants(mock_populate):
    """Editors + viewers; editors only; include_viewers=False."""
    editor_s = [_make_subject(1)]
    viewer_s = [_make_subject(2)]

    # Both editors and viewers
    mock_populate.side_effect = [editor_s, viewer_s]
    props: dict = {"editors": [1], "viewers": [2]}
    populate_subjects(props, exceptions := [])
    assert props["editors"] == editor_s
    assert props["viewers"] == viewer_s
    mock_populate.assert_any_call(
        [1],
        default_to_user=True,
        ensure_no_lockout=True,
        field_name="editors",
    )
    mock_populate.assert_any_call(
        [2],
        default_to_user=False,
        field_name="viewers",
    )
    assert not exceptions

    # Viewers absent from payload → not populated
    mock_populate.reset_mock(side_effect=True)
    mock_populate.return_value = editor_s
    props2: dict = {"editors": [1]}
    populate_subjects(props2, exceptions := [])
    assert "viewers" not in props2
    mock_populate.assert_called_once_with(
        [1],
        default_to_user=True,
        ensure_no_lockout=True,
        field_name="editors",
    )

    # include_viewers=False → viewers key untouched
    mock_populate.reset_mock(side_effect=True)
    mock_populate.return_value = editor_s
    props3: dict = {"editors": [1], "viewers": [2]}
    populate_subjects(props3, exceptions := [], include_viewers=False)
    assert props3["viewers"] == [2]  # unchanged raw IDs
    mock_populate.assert_called_once_with(
        [1],
        default_to_user=True,
        ensure_no_lockout=True,
        field_name="editors",
    )


@patch("superset.commands.utils.compute_subject_list")
def test_compute_subjects_all_variants(mock_compute):
    """Editors + viewers; no viewers attr; include_viewers=False."""
    model = MagicMock()
    model.editors = [_make_subject(1)]
    model.viewers = [_make_subject(2)]

    # Both editors and viewers
    mock_compute.side_effect = [[_make_subject(10)], [_make_subject(20)]]
    props: dict = {"editors": [10], "viewers": [20]}
    compute_subjects(model, props, exceptions := [])
    assert mock_compute.call_count == 2
    mock_compute.assert_any_call(
        model.editors,
        [10],
        ensure_no_lockout=True,
        field_name="editors",
    )
    mock_compute.assert_any_call(model.viewers, [20], field_name="viewers")
    assert not exceptions

    # Model without viewers attr → only editors
    mock_compute.reset_mock(side_effect=True)
    mock_compute.return_value = [_make_subject(10)]
    model_no_viewers = MagicMock(spec=["editors"])
    model_no_viewers.editors = [_make_subject(1)]
    compute_subjects(model_no_viewers, {"editors": [10]}, exceptions := [])
    mock_compute.assert_called_once_with(
        model_no_viewers.editors,
        [10],
        ensure_no_lockout=True,
        field_name="editors",
    )

    # include_viewers=False → skip viewers
    mock_compute.reset_mock(side_effect=True)
    mock_compute.return_value = [_make_subject(10)]
    compute_subjects(model, {"editors": [10]}, exceptions := [], include_viewers=False)
    mock_compute.assert_called_once_with(
        model.editors,
        [10],
        ensure_no_lockout=True,
        field_name="editors",
    )