superset2/tests/unit_tests/extensions/storage/test_persistent.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 persistent state storage implementation."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest
from flask import Flask

from superset.extensions.context import use_context
from superset.extensions.storage.persistent import (
    PersistentListOptions,
    PersistentSetOptions,
    PersistentState,
    SharedPersistentStateAccessor,
)
from superset.extensions.storage.persistent_dao import ExtensionStorageListEntry
from superset.utils import json
from tests.unit_tests.extensions.storage.conftest import (
    create_context,
    set_user,
)


@patch("superset.db")
def test_persistent_state_raises_without_context(
    mock_db: MagicMock, app: Flask
) -> None:
    """PersistentState operations raise RuntimeError without extension context."""
    with app.app_context():
        set_user(1)

        with pytest.raises(RuntimeError, match="within an extension context"):
            PersistentState.get("key")

        with pytest.raises(RuntimeError, match="within an extension context"):
            PersistentState.set("key", "value")

        with pytest.raises(RuntimeError, match="within an extension context"):
            PersistentState.remove("key")


def test_persistent_state_raises_without_user(app: Flask) -> None:
    """PersistentState operations raise RuntimeError without authenticated user."""
    ctx = create_context()

    with app.app_context(), use_context(ctx):
        with pytest.raises(RuntimeError, match="requires an authenticated user"):
            PersistentState.get("key")


@patch("superset.extensions.storage.persistent.ExtensionStorageDAO")
def test_persistent_state_get_returns_value(mock_dao: MagicMock, app: Flask) -> None:
    """PersistentState.get returns decoded value from DAO."""
    ctx = create_context()
    mock_dao.get_decoded_value.return_value = {"theme": "dark"}

    with app.app_context(), use_context(ctx):
        set_user(42)
        result = PersistentState.get("prefs")

    mock_dao.get_decoded_value.assert_called_once_with(
        "test-org.test-ext", "prefs", user_fk=42
    )
    assert result == {"theme": "dark"}


@patch("superset.extensions.storage.persistent.ExtensionStorageDAO")
def test_persistent_state_get_returns_none_when_missing(
    mock_dao: MagicMock, app: Flask
) -> None:
    """PersistentState.get returns None when key does not exist."""
    ctx = create_context()
    mock_dao.get_decoded_value.return_value = None

    with app.app_context(), use_context(ctx):
        set_user(42)
        result = PersistentState.get("missing")

    assert result is None


@patch("superset.db")
@patch("superset.extensions.storage.persistent.ExtensionStorageDAO")
def test_persistent_state_set_encodes_value(
    mock_dao: MagicMock, mock_db: MagicMock, app: Flask
) -> None:
    """PersistentState.set encodes value as JSON bytes."""
    ctx = create_context()

    with app.app_context(), use_context(ctx):
        set_user(42)
        PersistentState.set("prefs", {"theme": "dark"})

    expected_bytes = json.dumps({"theme": "dark"}).encode()
    mock_dao.set.assert_called_once_with(
        "test-org.test-ext",
        "prefs",
        expected_bytes,
        codec="json",
        user_fk=42,
        encrypt=False,
    )


@patch("superset.db")
@patch("superset.extensions.storage.persistent.ExtensionStorageDAO")
def test_persistent_state_set_with_encrypt_option(
    mock_dao: MagicMock, mock_db: MagicMock, app: Flask
) -> None:
    """PersistentState.set passes encrypt=True when options.encrypt is set."""
    ctx = create_context()

    with app.app_context(), use_context(ctx):
        set_user(42)
        PersistentState.set("token", "sk-...", PersistentSetOptions(encrypt=True))

    expected_bytes = json.dumps("sk-...").encode()
    mock_dao.set.assert_called_once_with(
        "test-org.test-ext",
        "token",
        expected_bytes,
        codec="json",
        user_fk=42,
        encrypt=True,
    )


@patch("superset.db")
@patch("superset.extensions.storage.persistent.ExtensionStorageDAO")
def test_shared_accessor_set_with_encrypt_option(
    mock_dao: MagicMock, mock_db: MagicMock, app: Flask
) -> None:
    """SharedPersistentStateAccessor.set passes encrypt=True when requested."""
    ctx = create_context()
    accessor = SharedPersistentStateAccessor()

    with app.app_context(), use_context(ctx):
        set_user(42)
        accessor.set("shared_token", "sk-...", PersistentSetOptions(encrypt=True))

    expected_bytes = json.dumps("sk-...").encode()
    mock_dao.set.assert_called_once_with(
        "test-org.test-ext",
        "shared_token",
        expected_bytes,
        codec="json",
        user_fk=None,
        encrypt=True,
    )


@patch("superset.db")
@patch("superset.extensions.storage.persistent.ExtensionStorageDAO")
def test_persistent_state_remove_deletes_entry(
    mock_dao: MagicMock, mock_db: MagicMock, app: Flask
) -> None:
    """PersistentState.remove calls DAO delete."""
    ctx = create_context()

    with app.app_context(), use_context(ctx):
        set_user(42)
        PersistentState.remove("prefs")

    mock_dao.delete_by_key.assert_called_once_with(
        "test-org.test-ext", "prefs", user_fk=42
    )


@patch("superset.extensions.storage.persistent.ExtensionStorageDAO")
def test_shared_accessor_uses_null_user_fk(mock_dao: MagicMock, app: Flask) -> None:
    """SharedPersistentStateAccessor uses user_fk=None for global scope."""
    ctx = create_context()
    mock_dao.get_decoded_value.return_value = "shared_value"

    accessor = SharedPersistentStateAccessor()

    with app.app_context(), use_context(ctx):
        set_user(42)
        result = accessor.get("config")

    mock_dao.get_decoded_value.assert_called_once_with(
        "test-org.test-ext", "config", user_fk=None
    )
    assert result == "shared_value"


@patch("superset.extensions.storage.persistent.ExtensionStorageDAO")
def test_persistent_state_list_decodes_entries_and_forwards_options(
    mock_dao: MagicMock, app: Flask
) -> None:
    """PersistentState.list decodes every entry's value (no SAFE_CODECS
    restriction for ambient backend code) and forwards the given options."""
    ctx = create_context()
    mock_dao.list_entries.return_value = (
        [
            ExtensionStorageListEntry(
                key="prefs",
                value=json.dumps({"theme": "dark"}).encode(),
                codec="json",
                is_encrypted=False,
            ),
        ],
        1,
    )

    with app.app_context(), use_context(ctx):
        set_user(42)
        result = PersistentState.list(PersistentListOptions(page=0, page_size=10))

    mock_dao.list_entries.assert_called_once_with(
        "test-org.test-ext",
        user_fk=42,
        resource_type=None,
        resource_uuid=None,
        page=0,
        page_size=10,
    )
    assert result.count == 1
    assert len(result.entries) == 1
    assert result.entries[0].key == "prefs"
    assert result.entries[0].value == {"theme": "dark"}
    assert result.entries[0].codec == "json"


@patch("superset.extensions.storage.persistent.ExtensionStorageDAO")
def test_persistent_state_list_forwards_custom_options(
    mock_dao: MagicMock, app: Flask
) -> None:
    """PersistentState.list forwards caller-supplied PersistentListOptions."""
    ctx = create_context()
    mock_dao.list_entries.return_value = ([], 0)

    with app.app_context(), use_context(ctx):
        set_user(42)
        PersistentState.list(
            PersistentListOptions(
                resource_type="dashboard",
                resource_uuid="uuid-1",
                page=2,
                page_size=25,
            )
        )

    mock_dao.list_entries.assert_called_once_with(
        "test-org.test-ext",
        user_fk=42,
        resource_type="dashboard",
        resource_uuid="uuid-1",
        page=2,
        page_size=25,
    )


@patch("superset.extensions.storage.persistent.ExtensionStorageDAO")
def test_persistent_state_list_handles_none_value(
    mock_dao: MagicMock, app: Flask
) -> None:
    """PersistentState.list does not attempt to decode a None value (e.g.
    a decryption failure), passing it through as None instead of raising."""
    ctx = create_context()
    mock_dao.list_entries.return_value = (
        [
            ExtensionStorageListEntry(
                key="broken",
                value=None,
                codec="json",
                is_encrypted=True,
            ),
        ],
        1,
    )

    with app.app_context(), use_context(ctx):
        set_user(42)
        result = PersistentState.list(PersistentListOptions(page=0, page_size=10))

    assert result.entries[0].value is None


@patch("superset.extensions.storage.persistent.ExtensionStorageDAO")
def test_shared_accessor_list_uses_null_user_fk(
    mock_dao: MagicMock, app: Flask
) -> None:
    """SharedPersistentStateAccessor.list uses user_fk=None for global scope."""
    ctx = create_context()
    mock_dao.list_entries.return_value = ([], 0)
    accessor = SharedPersistentStateAccessor()

    with app.app_context(), use_context(ctx):
        set_user(42)
        accessor.list(PersistentListOptions(page=0, page_size=10))

    mock_dao.list_entries.assert_called_once_with(
        "test-org.test-ext",
        user_fk=None,
        resource_type=None,
        resource_uuid=None,
        page=0,
        page_size=10,
    )