Spaces:

Leon4gr45
/

powerbi

Paused

powerbi / tests /unit_tests /sql /execution /conftest.py

Upload folder using huggingface_hub (part 17)

5b7ed07 verified 2 months ago

10.7 kB

# 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

	# 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