tests/unit_tests/sql/execution/conftest.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.

"""
Shared fixtures and helpers for SQL execution tests.

This module provides common mocks, fixtures, and helper functions used across
test_celery_task.py and test_executor.py to reduce code duplication.
"""

from contextlib import contextmanager
from typing import Any
from unittest.mock import MagicMock

import pandas as pd
import pytest
from flask import current_app
from pytest_mock import MockerFixture

from superset.common.db_query_status import QueryStatus as QueryStatusEnum
from superset.models.core import Database

# =============================================================================
# Core Fixtures
# =============================================================================


@pytest.fixture(autouse=True)
def mock_db_session(mocker: MockerFixture) -> MagicMock:
    """Mock database session for all tests to avoid foreign key constraints."""
    mock_session = MagicMock()
    mocker.patch("superset.sql.execution.executor.db.session", mock_session)
    mocker.patch("superset.sql.execution.celery_task.db.session", mock_session)
    return mock_session


@pytest.fixture
def mock_query() -> MagicMock:
    """Create a mock Query model."""
    query = MagicMock()
    query.id = 123
    query.database_id = 1
    query.sql = "SELECT * FROM users"
    query.status = QueryStatusEnum.PENDING
    query.error_message = None
    query.progress = 0
    query.end_time = None
    query.start_running_time = None
    query.executed_sql = None
    query.tmp_table_name = None
    query.catalog = None
    query.schema = "public"
    query.extra = {}
    query.set_extra_json_key = MagicMock()
    query.results_key = None
    query.select_as_cta = False
    query.rows = 0
    query.to_dict = MagicMock(return_value={"id": 123})
    query.database = MagicMock()
    query.database.db_engine_spec.extract_errors.return_value = []
    query.database.unique_name = "test_db"
    query.database.cache_timeout = 300
    return query


def _passthrough_mutate_sql_based_on_config(sql: str, **kwargs: Any) -> str:
    """Mirror the real `Database.mutate_sql_based_on_config` no-op default."""
    return sql


@pytest.fixture
def mock_database() -> MagicMock:
    """Create a mock Database."""
    database = MagicMock()
    database.id = 1
    database.unique_name = "test_db"
    database.cache_timeout = 300
    database.sqlalchemy_uri = "postgresql://localhost/test"
    database.db_engine_spec = MagicMock()
    database.db_engine_spec.engine = "postgresql"
    database.db_engine_spec.run_multiple_statements_as_one = False
    database.db_engine_spec.allows_sql_comments = True
    database.db_engine_spec.extract_errors = MagicMock(return_value=[])
    database.db_engine_spec.execute_with_cursor = MagicMock()
    database.db_engine_spec.get_cancel_query_id = MagicMock(return_value=None)
    database.db_engine_spec.patch = MagicMock()
    database.db_engine_spec.fetch_data = MagicMock(return_value=[])
    # Mirrors the real `Database.mutate_sql_based_on_config` default (no-op
    # when no `SQL_QUERY_MUTATOR` is configured), so SQL parsed from its
    # return value stays valid instead of an un-parseable `MagicMock`.
    database.mutate_sql_based_on_config = MagicMock(
        side_effect=_passthrough_mutate_sql_based_on_config
    )
    return database


@pytest.fixture
def mock_result_set() -> MagicMock:
    """Create a mock SupersetResultSet."""
    result_set = MagicMock()
    result_set.size = 2
    result_set.columns = [{"name": "id"}, {"name": "name"}]
    result_set.pa_table = MagicMock()
    result_set.to_pandas_df = MagicMock(
        return_value=pd.DataFrame({"id": [1, 2], "name": ["Alice", "Bob"]})
    )
    return result_set


@pytest.fixture
def database() -> Database:
    """Create a real test database instance."""
    return Database(
        id=1,
        database_name="test_db",
        sqlalchemy_uri="sqlite://",
        allow_dml=False,
    )


@pytest.fixture
def database_with_dml() -> Database:
    """Create a real test database instance with DML allowed."""
    return Database(
        id=2,
        database_name="test_db_dml",
        sqlalchemy_uri="sqlite://",
        allow_dml=True,
    )


# =============================================================================
# Helper Functions for Mock Creation
# =============================================================================


def create_mock_cursor(
    column_names: list[str],
    data: list[tuple[Any, ...]] | None = None,
) -> MagicMock:
    """
    Create a mock database cursor with column description.

    Args:
        column_names: List of column names
        data: Optional data to return from fetchall()

    Returns:
        Configured MagicMock cursor
    """
    mock_cursor = MagicMock()
    mock_cursor.description = [
        (name, None, None, None, None, None, None) for name in column_names
    ]
    if data is not None:
        mock_cursor.fetchall.return_value = data
    return mock_cursor


def create_mock_connection(mock_cursor: MagicMock | None = None) -> MagicMock:
    """
    Create a mock database connection.

    Args:
        mock_cursor: Optional cursor to return from cursor()

    Returns:
        Configured MagicMock connection with context manager support
    """
    if mock_cursor is None:
        mock_cursor = create_mock_cursor([])

    mock_conn = MagicMock()
    mock_conn.cursor.return_value = mock_cursor
    mock_conn.close = MagicMock()
    mock_conn.__enter__ = MagicMock(return_value=mock_conn)
    mock_conn.__exit__ = MagicMock(return_value=False)
    return mock_conn


def setup_mock_raw_connection(
    mock_database: MagicMock,
    mock_connection: MagicMock | None = None,
) -> MagicMock:
    """
    Setup get_raw_connection as a context manager on a mock database.

    Args:
        mock_database: The database mock to configure
        mock_connection: Optional connection to yield

    Returns:
        The configured mock connection
    """
    if mock_connection is None:
        mock_connection = create_mock_connection()

    @contextmanager
    def _raw_connection(
        catalog: str | None = None,
        schema: str | None = None,
        nullpool: bool = True,
        source: Any | None = None,
    ):
        yield mock_connection

    mock_database.get_raw_connection = _raw_connection
    return mock_connection


def setup_db_session_query_mock(
    mock_db_session: MagicMock,
    return_value: Any = None,
) -> None:
    """
    Setup database session query chain for query lookup.

    Args:
        mock_db_session: The database session mock
        return_value: Value to return from one_or_none()
    """
    filter_mock = mock_db_session.query.return_value.filter_by.return_value
    filter_mock.one_or_none.return_value = return_value


def mock_query_execution(
    mocker: MockerFixture,
    database: Database,
    return_data: list[tuple[Any, ...]],
    column_names: list[str],
) -> MagicMock:
    """
    Mock the raw connection execution path for testing.

    This helper sets up all necessary mocks for executing a query through
    the database engine spec and returning results.

    Args:
        mocker: pytest-mock fixture
        database: Database instance to mock
        return_data: Data to return from fetch_data, e.g. [(1, "Alice"), (2, "Bob")]
        column_names: Column names for the result, e.g. ["id", "name"]

    Returns:
        The mock for get_raw_connection, so tests can make assertions on it
    """
    from superset.result_set import SupersetResultSet

    # Mock cursor and connection
    mock_cursor = create_mock_cursor(column_names, return_data)
    mock_conn = create_mock_connection(mock_cursor)

    get_raw_conn_mock = mocker.patch.object(
        database, "get_raw_connection", return_value=mock_conn
    )
    mocker.patch.object(
        database, "mutate_sql_based_on_config", side_effect=lambda sql, **kw: sql
    )
    mocker.patch.object(database.db_engine_spec, "execute")
    mocker.patch.object(database.db_engine_spec, "fetch_data", return_value=return_data)

    # Create a real SupersetResultSet that converts to DataFrame properly
    mock_result_set = MagicMock(spec=SupersetResultSet)
    mock_result_set.to_pandas_df.return_value = pd.DataFrame(
        return_data, columns=column_names
    )
    mocker.patch("superset.result_set.SupersetResultSet", return_value=mock_result_set)

    return get_raw_conn_mock


# =============================================================================
# Composite Fixtures for Common Test Patterns
# =============================================================================


@pytest.fixture
def default_sql_config(mocker: MockerFixture) -> None:
    """Patch app config with default SQL execution settings."""
    mocker.patch.dict(
        current_app.config,
        {
            "SQL_QUERY_MUTATOR": None,
            "SQLLAB_TIMEOUT": 30,
            "SQL_MAX_ROW": None,
            "QUERY_LOGGER": None,
        },
    )


@pytest.fixture
def mock_celery_task(mocker: MockerFixture) -> MagicMock:
    """Mock the Celery task for SQL execution."""
    return mocker.patch("superset.sql.execution.celery_task.execute_sql_task")


def setup_cache_mocks(
    mocker: MockerFixture,
    get_result: Any = None,
    store_side_effect: Any = None,
) -> tuple[MagicMock, MagicMock]:
    """
    Setup cache get/store mocks for executor tests.

    Args:
        mocker: pytest-mock fixture
        get_result: Value to return from _get_from_cache
        store_side_effect: Optional side effect for _store_in_cache

    Returns:
        Tuple of (mock_get, mock_store)
    """
    from superset.sql.execution.executor import SQLExecutor

    mock_get = mocker.patch.object(
        SQLExecutor, "_get_from_cache", return_value=get_result
    )
    mock_store = mocker.patch.object(
        SQLExecutor, "_store_in_cache", side_effect=store_side_effect
    )
    return mock_get, mock_store