pytest-spark

Warnings

• gotcha Configuring `SPARK_HOME` has multiple methods (environment variable, `pytest.ini`, `--spark_home` CLI option) which are read in a specific order (CLI > `pytest.ini` > ENV). If `pyspark` is installed via pip, setting `SPARK_HOME` might not be necessary, leading to confusion or unexpected behavior if mismatched.
  Fix: Prefer installing `pyspark` via `pip` and omit `SPARK_HOME` if possible. If explicit `SPARK_HOME` is needed, use `pytest.ini` for project-level consistency or `--spark_home` for specific runs, understanding their precedence.
• breaking The `spark_context` fixture is not supported when using Spark Connect functionality. If you're working with Spark 3.4+ and Spark Connect, you must use the `spark_session` fixture instead.
  Fix: Use the `spark_session` fixture for tests involving Spark Connect.
• gotcha By default, the `spark_session` fixture creates a SparkSession with Hive support enabled. If Hive jars are not desired or cause conflicts, you can explicitly disable Hive support by adding `spark_options = spark.sql.catalogImplementation: in-memory` to your `pytest.ini`.
  Fix: Add `spark_options = spark.sql.catalogImplementation: in-memory` under the `[pytest]` section in your `pytest.ini` to explicitly disable Hive support.
• gotcha Comparing Spark DataFrames for equality can be challenging directly due to potential differences in row order, column order, or schema. Direct `==` comparison often fails even for logically identical DataFrames.
  Fix: For robust DataFrame comparison, convert both expected and actual DataFrames to Pandas DataFrames, sort them by common keys (if order doesn't matter), and then use `pandas.testing.assert_frame_equal(df1.toPandas().sort_values(...), df2.toPandas().sort_values(...), check_like=True)` to ignore column order. Libraries like `chispa` also provide Spark DataFrame equality assertions.
• gotcha When running tests in parallel with `pytest-xdist`, session-scoped Spark fixtures (like `spark_session`) can interfere with each other if not properly isolated. Each parallel process might try to use the same temporary directories or resources, leading to data races or failures.
  Fix: Utilize `pytest`'s `tmp_path_factory` fixture to generate a unique temporary directory for each `pytest-xdist` worker process, and configure Spark to use these isolated directories for its local storage (e.g., `spark.local.dir`).

Install

• pip install pytest-spark Install stable version

Imports

• spark_session

  def test_example(spark_session):

  pytest-spark provides fixtures that pytest auto-discovers; no direct import from 'pytest_spark' itself is typically needed in test files.
• spark_context

  def test_example(spark_context):

  pytest-spark provides fixtures that pytest auto-discovers; no direct import from 'pytest_spark' itself is typically needed in test files.

Quickstart

To get started, define a `spark_session` fixture in a `conftest.py` file. This fixture will automatically provide a SparkSession to your tests. Then, write tests that accept `spark_session` as an argument. You can run tests using `pytest` from your terminal.

import pytest
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, IntegerType

# conftest.py (placed in your project's root or tests directory)
@pytest.fixture(scope="session")
def spark_session():
    """
    Fixture for creating a SparkSession for testing.
    This SparkSession is reused across all tests in the session.
    """
    spark = SparkSession.builder \
        .master("local[*]") \
        .appName("pytest-spark-session") \
        .config("spark.driver.memory", "2g") \
        .getOrCreate()
    yield spark
    spark.stop()

# test_example.py (a sample test file)
def test_data_frame_creation(spark_session):
    schema = StructType([
        StructField("name", StringType(), True),
        StructField("age", IntegerType(), True)
    ])
    data = [("Alice", 1), ("Bob", 2)]
    df = spark_session.createDataFrame(data, schema)
    
    assert df.count() == 2
    assert df.columns == ["name", "age"]
    assert df.collect()[0].name == "Alice"