sparkless-testing 0.1.0__py3-none-any.whl

sparkless_testing/protocols.py

@@ -0,0 +1,145 @@
+"""
+Protocol definitions for engine-agnostic Spark interfaces.
+
+These protocols represent the minimal surface area used by sparkless-testing.
+Any engine (PySpark, sparkless, or other) must satisfy these to work with
+the testing framework.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class ColumnProtocol(Protocol):
+    """Column-like object with comparison and basic ops."""
+
+    def isNotNull(self) -> Any: ...
+    def isNull(self) -> Any: ...
+    def cast(self, dataType: Any) -> Any: ...
+    def __eq__(self, other: Any) -> Any: ...
+    def __ne__(self, other: Any) -> Any: ...
+    def __gt__(self, other: Any) -> Any: ...
+    def __lt__(self, other: Any) -> Any: ...
+    def __ge__(self, other: Any) -> Any: ...
+    def __le__(self, other: Any) -> Any: ...
+    def __add__(self, other: Any) -> Any: ...
+    def __sub__(self, other: Any) -> Any: ...
+    def __mul__(self, other: Any) -> Any: ...
+    def __truediv__(self, other: Any) -> Any: ...
+
+
+@runtime_checkable
+class DataFrameProtocol(Protocol):
+    """DataFrame-like object used throughout testing."""
+
+    # Core accessors
+    def schema(self) -> Any: ...
+    @property
+    def columns(self) -> list[str]: ...
+
+    # Actions
+    def count(self) -> int: ...
+    def collect(self) -> list[Any]: ...
+    def show(self, n: int = 20, truncate: bool = True) -> None: ...
+
+    # Transformations
+    def filter(self, condition: Any) -> Any: ...
+    def select(self, *cols: Any) -> Any: ...
+    def withColumn(self, colName: str, col: Any) -> Any: ...
+    def withColumnRenamed(self, existing: str, new: str) -> Any: ...
+    def groupBy(self, *cols: Any) -> Any: ...
+    def agg(self, *exprs: Any, **kwargs: Any) -> Any: ...
+    def limit(self, num: int) -> Any: ...
+    def cache(self) -> Any: ...
+
+    # SQL helpers
+    def createOrReplaceTempView(self, name: str) -> None: ...
+
+    # Writer
+    @property
+    def write(self) -> Any: ...
+
+
+@runtime_checkable
+class FunctionsProtocol(Protocol):
+    """Functions module interface (col, lit, aggregations, etc.)."""
+
+    def col(self, col_name: str) -> ColumnProtocol: ...
+    def expr(self, expr: str) -> ColumnProtocol: ...
+    def lit(self, value: Any) -> ColumnProtocol: ...
+    def when(self, condition: ColumnProtocol, value: Any) -> ColumnProtocol: ...
+    def count(self, col: Any = "*") -> ColumnProtocol: ...
+    def countDistinct(self, *cols: Any) -> ColumnProtocol: ...
+    def sum(self, col: Any) -> ColumnProtocol: ...
+    def max(self, col: Any) -> ColumnProtocol: ...
+    def min(self, col: Any) -> ColumnProtocol: ...
+    def avg(self, col: Any) -> ColumnProtocol: ...
+    def length(self, col: Any) -> ColumnProtocol: ...
+    def date_trunc(self, fmt: str, col: Any) -> ColumnProtocol: ...
+    def dayofweek(self, col: Any) -> ColumnProtocol: ...
+    def current_timestamp(self) -> ColumnProtocol: ...
+
+
+@runtime_checkable
+class TypesProtocol(Protocol):
+    """Types namespace used for schemas and fields."""
+
+    StructType: Any
+    StructField: Any
+    StringType: Any
+    IntegerType: Any
+    FloatType: Any
+    DoubleType: Any
+    LongType: Any
+    TimestampType: Any
+    BooleanType: Any
+
+
+@runtime_checkable
+class WindowProtocol(Protocol):
+    """Window spec placeholder."""
+
+    def orderBy(self, *cols: Any, **kwargs: Any) -> Any: ...
+    def partitionBy(self, *cols: Any) -> Any: ...
+
+
+@runtime_checkable
+class AnalysisExceptionProtocol(Protocol):
+    """Exception type placeholder for analysis errors."""
+
+    @property
+    def desc(self) -> str: ...
+
+
+@runtime_checkable
+class SparkSessionProtocol(Protocol):
+    """SparkSession-like interface."""
+
+    @property
+    def catalog(self) -> Any: ...
+    @property
+    def conf(self) -> Any: ...
+
+    def table(self, name: str) -> DataFrameProtocol: ...
+    def createDataFrame(
+        self, data: Any, schema: Any = None, samplingRatio: Any = None
+    ) -> DataFrameProtocol: ...
+    def sql(self, sqlQuery: str) -> DataFrameProtocol: ...
+    def stop(self) -> None: ...
+
+    # Builder/config checks
+    @property
+    def _jsparkSession(self) -> Any: ...  # optional; used for ids
+
+
+__all__ = [
+    "ColumnProtocol",
+    "DataFrameProtocol",
+    "FunctionsProtocol",
+    "TypesProtocol",
+    "WindowProtocol",
+    "AnalysisExceptionProtocol",
+    "SparkSessionProtocol",
+]

sparkless_testing/pytest_fixtures.py

@@ -0,0 +1,235 @@
+"""
+Pytest fixtures for sparkless-testing.
+
+This module provides pytest fixtures for easy test integration with both
+sparkless (mock) and PySpark (real) engines.
+"""
+
+from __future__ import annotations
+
+import os
+import uuid
+from collections.abc import Generator
+from typing import Any
+
+import pytest
+
+from .engine import EngineType, auto_configure_engine, get_engine
+from .session import SessionConfig, create_mock_session, create_pyspark_session, create_session
+from .utils import get_spark_mode
+
+
+def _get_engine_type_from_marker(request: pytest.FixtureRequest) -> EngineType | None:
+    """Get engine type from pytest marker if present."""
+    marker = request.node.get_closest_marker("spark_engine")
+    if marker:
+        engine_arg = marker.args[0] if marker.args else None
+        if engine_arg:
+            engine_str = str(engine_arg).lower()
+            if engine_str == "mock":
+                return EngineType.MOCK
+            elif engine_str == "pyspark":
+                return EngineType.PYSPARK
+    return None
+
+
+def _determine_engine_type(request: pytest.FixtureRequest) -> EngineType:
+    """Determine which engine type to use based on markers and environment."""
+    # Check pytest marker first
+    marker_engine = _get_engine_type_from_marker(request)
+    if marker_engine:
+        return marker_engine
+
+    # Check environment variable
+    spark_mode = get_spark_mode()
+    if spark_mode == "mock":
+        return EngineType.MOCK
+    elif spark_mode == "pyspark":
+        return EngineType.PYSPARK
+    else:
+        # Auto-detect
+        return EngineType.AUTO
+
+
+@pytest.fixture(scope="function")  # type: ignore[untyped-decorator]
+def spark_session(request: pytest.FixtureRequest) -> Generator[Any, None, None]:
+    """
+    Main fixture that provides appropriate SparkSession based on configuration.
+
+    This fixture respects:
+    - `SPARK_MODE` environment variable: `mock`, `pyspark`, or `auto` (default)
+    - `@pytest.mark.spark_engine("mock")` or `@pytest.mark.spark_engine("pyspark")` markers
+    - Auto-detection if only one engine is available
+    - Parametrization via `@pytest.mark.parametrize_engines`
+
+    Example:
+        def test_my_function(spark_session):
+            df = spark_session.createDataFrame([{"id": 1}])
+            assert df.count() == 1
+    """
+    # Check if parametrized (for parametrize_engines support)
+    param = getattr(request, "param", None)
+    if param is not None:
+        # Use the parametrized engine type
+        engine_type = param
+    else:
+        # Determine engine type from markers/environment
+        engine_type = _determine_engine_type(request)
+
+    # Auto-configure engine if not already configured
+    try:
+        get_engine()
+    except RuntimeError:
+        if engine_type == EngineType.AUTO:
+            auto_configure_engine()
+        else:
+            auto_configure_engine(engine_type)
+
+    # Create session with Delta Lake disabled for parallel testing compatibility
+    # Use process ID and UUID for unique session names (pytest-xdist uses separate processes)
+    unique_id = f"{os.getpid()}-{uuid.uuid4().hex[:8]}"
+    config = SessionConfig(app_name=f"test-{unique_id}", enable_delta=False)
+    spark = create_session(engine_type=engine_type, config=config)
+
+    yield spark
+
+    # Cleanup
+    try:
+        spark.stop()
+    except Exception:
+        pass
+
+
+@pytest.fixture(scope="function")  # type: ignore[untyped-decorator]
+def mock_spark_session() -> Generator[Any, None, None]:
+    """
+    Fixture that explicitly uses sparkless (mock) engine.
+
+    Example:
+        def test_mock_only(mock_spark_session):
+            # Only runs with sparkless
+            df = mock_spark_session.createDataFrame([{"id": 1}])
+            assert df.count() == 1
+    """
+    # Auto-configure mock engine if not already configured
+    try:
+        get_engine()
+    except RuntimeError:
+        auto_configure_engine(EngineType.MOCK)
+
+    spark = create_mock_session()
+
+    yield spark
+
+    # Cleanup
+    try:
+        spark.stop()
+    except Exception:
+        pass
+
+
+@pytest.fixture(scope="function")  # type: ignore[untyped-decorator]
+def pyspark_session() -> Generator[Any, None, None]:
+    """
+    Fixture that explicitly uses PySpark (real) engine.
+
+    Example:
+        def test_pyspark_only(pyspark_session):
+            # Only runs with PySpark
+            df = pyspark_session.createDataFrame([{"id": 1}])
+            assert df.count() == 1
+    """
+    # Auto-configure PySpark engine if not already configured
+    try:
+        get_engine()
+    except RuntimeError:
+        auto_configure_engine(EngineType.PYSPARK)
+
+    # Disable Delta Lake for parallel testing compatibility
+    unique_id = f"{os.getpid()}-{uuid.uuid4().hex[:8]}"
+    config = SessionConfig(app_name=f"test-{unique_id}", enable_delta=False)
+    spark = create_pyspark_session(config=config)
+
+    yield spark
+
+    # Cleanup
+    try:
+        spark.stop()
+    except Exception:
+        pass
+
+
+@pytest.fixture(scope="function")  # type: ignore[untyped-decorator]
+def spark_functions() -> Any:
+    """
+    Fixture that provides the functions module (F) for the current engine.
+
+    **Important**: For PySpark, functions like `F.col()` require an active SparkContext.
+    Always use this fixture together with `spark_session` to ensure the session is active.
+
+    Example:
+        def test_with_functions(spark_session, spark_functions):
+            F = spark_functions
+            df = spark_session.createDataFrame([{"name": "Alice"}])
+            result = df.select(F.upper(F.col("name"))).collect()
+            assert result[0][0] == "ALICE"
+
+    Note:
+        If you use this fixture without `spark_session`, PySpark functions may fail
+        with an AssertionError if no SparkContext is active. Always use both fixtures together.
+    """
+    try:
+        return get_engine().functions
+    except RuntimeError:
+        # Auto-configure if not configured
+        auto_configure_engine()
+        return get_engine().functions
+
+
+@pytest.fixture(scope="function")  # type: ignore[untyped-decorator]
+def spark_types() -> Any:
+    """
+    Fixture that provides the types module for the current engine.
+
+    Example:
+        def test_with_types(spark_session, spark_types):
+            from spark_types import StructType, StructField, StringType
+            schema = StructType([StructField("name", StringType())])
+            df = spark_session.createDataFrame([{"name": "Alice"}], schema)
+    """
+    try:
+        return get_engine().types
+    except RuntimeError:
+        # Auto-configure if not configured
+        auto_configure_engine()
+        return get_engine().types
+
+
+@pytest.fixture(scope="function")  # type: ignore[untyped-decorator]
+def spark_engine_type() -> str:
+    """
+    Fixture that provides the current engine type as a string.
+
+    Returns:
+        'mock' or 'pyspark'
+
+    Example:
+        def test_engine_type(spark_engine_type):
+            assert spark_engine_type in ("mock", "pyspark")
+    """
+    try:
+        return get_engine().engine_name
+    except RuntimeError:
+        # Auto-configure if not configured
+        auto_configure_engine()
+        return get_engine().engine_name
+
+
+__all__ = [
+    "spark_session",
+    "mock_spark_session",
+    "pyspark_session",
+    "spark_functions",
+    "spark_types",
+    "spark_engine_type",
+]

sparkless_testing/session.py

@@ -0,0 +1,262 @@
+"""
+Session factory for creating SparkSession instances.
+
+This module provides functions to create SparkSession instances for both
+sparkless (mock) and PySpark (real) engines with proper configuration.
+"""
+
+from __future__ import annotations
+
+import os
+import uuid
+import warnings
+from dataclasses import dataclass
+from typing import Any
+
+from .engine import EngineType, auto_configure_engine, get_engine
+from .utils import get_app_name_prefix, get_warehouse_dir
+
+
+@dataclass
+class SessionConfig:
+    """Configuration for SparkSession creation."""
+
+    app_name: str
+    warehouse_dir: str | None = None
+    enable_delta: bool = True
+    enable_ui: bool = False
+    shuffle_partitions: int = 1
+    parallelism: int = 1
+    adaptive_enabled: bool = False
+
+
+def create_mock_session(
+    app_name: str | None = None,
+    config: SessionConfig | None = None,
+) -> Any:
+    """
+    Create a sparkless (mock) SparkSession.
+
+    Args:
+        app_name: Application name (defaults to auto-generated)
+        config: Optional SessionConfig (currently unused for mock sessions)
+
+    Returns:
+        SparkSession instance from sparkless
+    """
+    try:
+        from sparkless import SparkSession
+    except ImportError as e:
+        raise RuntimeError(
+            "sparkless is not installed. Install it with: pip install sparkless"
+        ) from e
+
+    # Auto-configure engine if not already configured
+    try:
+        get_engine()
+    except RuntimeError:
+        auto_configure_engine(EngineType.MOCK)
+
+    # Generate app name if not provided
+    if app_name is None:
+        app_name = f"{get_app_name_prefix()}-mock-{os.getpid()}"
+
+    # Create mock Spark session
+    spark = SparkSession(app_name)
+
+    # Create test schema if possible
+    try:
+        spark.sql("CREATE SCHEMA IF NOT EXISTS test_schema")
+    except Exception:
+        pass  # Ignore schema creation errors
+
+    return spark
+
+
+def create_pyspark_session(
+    app_name: str | None = None,
+    config: SessionConfig | None = None,
+) -> Any:
+    """
+    Create a PySpark SparkSession with optional Delta Lake support.
+
+    Args:
+        app_name: Application name (defaults to auto-generated)
+        config: Optional SessionConfig for session customization
+
+    Returns:
+        SparkSession instance from PySpark
+    """
+    try:
+        from pyspark.sql import SparkSession
+    except ImportError as e:
+        raise RuntimeError("PySpark is not installed. Install it with: pip install pyspark") from e
+
+    # Auto-configure engine if not already configured
+    try:
+        get_engine()
+    except RuntimeError:
+        auto_configure_engine(EngineType.PYSPARK)
+
+    # Skip environment variable manipulation for parallel testing compatibility
+    # PySpark will use the current Python interpreter by default
+    # Setting these can cause contention in parallel execution
+
+    # Get worker ID for concurrent testing isolation (pytest-xdist)
+    worker_id = os.environ.get("PYTEST_XDIST_WORKER", "gw0")
+
+    # Generate app name if not provided
+    # Priority: explicit app_name parameter > config.app_name > auto-generated
+    if app_name is None:
+        if config and config.app_name:
+            app_name = config.app_name
+        else:
+            app_name = f"{get_app_name_prefix()}-pyspark-{worker_id}-{uuid.uuid4().hex[:8]}"
+
+    # Determine warehouse directory with unique identifier per process/test
+    warehouse_dir = None
+    if config and config.warehouse_dir:
+        warehouse_dir = config.warehouse_dir
+    else:
+        # Add unique identifier to avoid collisions across parallel workers
+        unique_id = uuid.uuid4().hex[:8]
+        base_dir = get_warehouse_dir() or f"/tmp/spark-warehouse-{os.getpid()}-{unique_id}"
+        warehouse_dir = base_dir
+
+    # Build Spark session builder - use local[1] to ensure single-threaded execution
+    # This prevents resource contention in parallel test execution
+    builder = (
+        SparkSession.builder.appName(app_name)
+        .master("local[1]")  # Explicitly use 1 thread, not local[*]
+        .config("spark.sql.warehouse.dir", warehouse_dir)
+        .config("spark.driver.host", "127.0.0.1")
+        .config("spark.driver.bindAddress", "127.0.0.1")
+        .config("spark.executor.cores", "1")  # Ensure only 1 core per executor
+        .config("spark.executor.instances", "1")  # Ensure only 1 executor instance
+    )
+
+    # Only add optional configs if explicitly provided
+    if config:
+        if config.enable_ui:
+            builder = builder.config("spark.ui.enabled", "true")
+        if config.shuffle_partitions != 1:
+            builder = builder.config("spark.sql.shuffle.partitions", str(config.shuffle_partitions))
+        if config.parallelism != 1:
+            builder = builder.config("spark.default.parallelism", str(config.parallelism))
+        if config.adaptive_enabled:
+            builder = builder.config("spark.sql.adaptive.enabled", "true")
+
+    # Add Delta Lake support if enabled and available
+    enable_delta = config.enable_delta if config else True
+    if enable_delta:
+        try:
+            from delta import configure_spark_with_delta_pip
+
+            builder = builder.config(
+                "spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension"
+            ).config(
+                "spark.sql.catalog.spark_catalog",
+                "org.apache.spark.sql.delta.catalog.DeltaCatalog",
+            )
+            builder = configure_spark_with_delta_pip(builder)
+        except ImportError:
+            # Delta Lake not available, continue without it
+            pass
+
+    # Stop any existing session in this process first to ensure clean state
+    # This prevents conflicts when multiple tests run in the same process
+    try:
+        existing = SparkSession.getActiveSession()
+        if existing is not None:
+            existing.stop()
+    except Exception:
+        pass
+
+    # Use getOrCreate() with unique app name - PySpark will create a new session
+    # for each unique app name, ensuring test isolation
+    spark = builder.getOrCreate()
+
+    # Minimal post-creation setup - avoid operations that could block
+    # Set log level to WARN to reduce noise (non-blocking)
+    try:
+        if hasattr(spark, "sparkContext") and spark.sparkContext:
+            spark.sparkContext.setLogLevel("WARN")
+    except Exception:
+        # Ignore errors - log level setting is not critical
+        pass
+
+    # No environment variable restoration needed since we're not modifying them
+
+    return spark
+
+
+def create_session(
+    engine_type: EngineType | None = None,
+    app_name: str | None = None,
+    config: SessionConfig | None = None,
+) -> Any:
+    """
+    Factory function that creates appropriate session based on engine type.
+
+    Args:
+        engine_type: Engine type to use (defaults to auto-detect)
+        app_name: Application name (defaults to auto-generated)
+        config: Optional SessionConfig for session customization
+
+    Returns:
+        SparkSession instance
+    """
+    # Auto-detect engine type if not specified
+    if engine_type is None or engine_type == EngineType.AUTO:
+        from .engine import detect_available_engines
+
+        available = detect_available_engines()
+        if EngineType.MOCK in available:
+            engine_type = EngineType.MOCK
+        elif EngineType.PYSPARK in available:
+            engine_type = EngineType.PYSPARK
+        else:
+            raise RuntimeError(
+                "No Spark engines available. Install either 'sparkless' or 'pyspark'."
+            )
+
+    # Create session based on engine type
+    if engine_type == EngineType.MOCK:
+        return create_mock_session(app_name, config)
+    elif engine_type == EngineType.PYSPARK:
+        return create_pyspark_session(app_name, config)
+    else:
+        raise ValueError(f"Unknown engine type: {engine_type}")
+
+
+def stop_all_sessions() -> None:
+    """
+    Stop all active Spark sessions.
+
+    This is a cleanup utility for test teardown.
+    """
+    # Try to stop PySpark sessions
+    try:
+        from pyspark.sql import SparkSession
+
+        active_session = SparkSession.getActiveSession()
+        if active_session:
+            active_session.stop()
+    except Exception as e:
+        warnings.warn(
+            f"Failed to stop existing PySpark session during cleanup: {e}",
+            UserWarning,
+            stacklevel=2,
+        )
+
+    # Note: sparkless sessions don't need explicit stopping,
+    # but we could add cleanup here if needed
+
+
+__all__ = [
+    "SessionConfig",
+    "create_mock_session",
+    "create_pyspark_session",
+    "create_session",
+    "stop_all_sessions",
+]