Testing

PySpark Testing with pytest: Unit Tests for DataFrame Transformations

HelpMeTest

23 Aug 2026 — 5 min read

PySpark code is hard to test because most engineers treat Spark jobs as scripts — a chain of transformations from input to output with no unit boundaries. The result is code where bugs only surface when you run the full job against real data, which takes 20 minutes and costs money.

The fix is straightforward: write testable PySpark by separating transformation logic from Spark session management, then test transformations as pure functions that take DataFrames and return DataFrames. This post covers the patterns that make PySpark code actually testable.

The Core Problem: Coupling to SparkSession

# Bad: SparkSession is global, transformations depend on it implicitly
spark = SparkSession.builder.appName("ETL").getOrCreate()

def process_orders():
    df = spark.read.parquet("s3://bucket/orders/")  # hard to mock
    return df.filter(df.status == "completed").groupBy("customer_id").sum("amount")

# Good: functions take and return DataFrames
def filter_completed_orders(df: DataFrame) -> DataFrame:
    return df.filter(df.status == "completed")

def compute_customer_revenue(df: DataFrame) -> DataFrame:
    return df.groupBy("customer_id").agg(
        F.sum("amount").alias("total_revenue"),
        F.count("*").alias("order_count")
    )

Now filter_completed_orders and compute_customer_revenue are pure functions — testable without touching S3 or a real cluster.

Setting Up pytest with PySpark

pip install pyspark pytest chispa

chispa is a library for asserting DataFrame equality in pytest with useful diff output.

# conftest.py
import pytest
from pyspark.sql import SparkSession

@pytest.fixture(scope="session")
def spark():
    """Shared SparkSession for all tests. Created once, reused across session."""
    spark = (SparkSession.builder
             .master("local[2]")
             .appName("unit-tests")
             .config("spark.sql.shuffle.partitions", "2")  # fast for small test data
             .config("spark.default.parallelism", "2")
             .config("spark.sql.adaptive.enabled", "false")
             .getOrCreate())
    yield spark
    spark.stop()

Testing DataFrame Transformations

# tests/test_order_transforms.py
import pytest
from pyspark.sql import Row
from chispa.dataframe_comparer import assert_df_equality
from myetl.transforms.orders import filter_completed_orders, compute_customer_revenue

def test_filter_completed_orders_keeps_only_completed(spark):
    input_df = spark.createDataFrame([
        Row(order_id=1, customer_id=10, amount=100.0, status="completed"),
        Row(order_id=2, customer_id=11, amount=50.0,  status="pending"),
        Row(order_id=3, customer_id=10, amount=75.0,  status="cancelled"),
        Row(order_id=4, customer_id=12, amount=200.0, status="completed"),
    ])
    
    result = filter_completed_orders(input_df)
    
    assert result.count() == 2
    assert result.filter("status != 'completed'").count() == 0

def test_compute_customer_revenue_aggregates_correctly(spark):
    input_df = spark.createDataFrame([
        Row(order_id=1, customer_id=10, amount=100.0, status="completed"),
        Row(order_id=2, customer_id=10, amount=50.0,  status="completed"),
        Row(order_id=3, customer_id=11, amount=200.0, status="completed"),
    ])
    
    expected = spark.createDataFrame([
        Row(customer_id=10, total_revenue=150.0, order_count=2),
        Row(customer_id=11, total_revenue=200.0, order_count=1),
    ])
    
    result = compute_customer_revenue(input_df)
    
    assert_df_equality(result, expected, ignore_row_order=True)

def test_compute_customer_revenue_handles_empty_input(spark):
    from pyspark.sql.types import StructType, StructField, IntegerType, DoubleType
    
    schema = StructType([
        StructField("order_id", IntegerType()),
        StructField("customer_id", IntegerType()),
        StructField("amount", DoubleType()),
        StructField("status", StringType()),
    ])
    
    empty_df = spark.createDataFrame([], schema)
    result = compute_customer_revenue(empty_df)
    
    assert result.count() == 0

Schema Validation Tests

Catching schema changes early prevents runtime failures:

# tests/test_schema_validation.py
from pyspark.sql.types import (StructType, StructField, 
                                StringType, DoubleType, LongType, TimestampType)

EXPECTED_ORDERS_SCHEMA = StructType([
    StructField("order_id",    LongType(),      nullable=False),
    StructField("customer_id", LongType(),      nullable=False),
    StructField("amount",      DoubleType(),    nullable=True),
    StructField("status",      StringType(),    nullable=False),
    StructField("created_at",  TimestampType(), nullable=False),
])

def test_orders_schema_matches_contract(spark):
    """Schema changes should be explicit — this test forces that."""
    df = spark.read.parquet("tests/fixtures/orders_sample.parquet")
    
    for field in EXPECTED_ORDERS_SCHEMA.fields:
        actual_field = df.schema[field.name]
        assert actual_field.dataType == field.dataType, \
            f"Column '{field.name}': expected {field.dataType}, got {actual_field.dataType}"
        assert actual_field.nullable == field.nullable, \
            f"Column '{field.name}' nullability changed"

def test_transform_output_schema_is_stable(spark):
    """Transformation output schema should match documented contract."""
    input_df = spark.createDataFrame([
        Row(order_id=1, customer_id=10, amount=100.0, status="completed",
            created_at=datetime(2026, 1, 1))
    ])
    
    result = compute_customer_revenue(input_df)
    
    assert "customer_id" in result.columns
    assert "total_revenue" in result.columns
    assert "order_count" in result.columns
    # No unexpected extra columns
    assert set(result.columns) == {"customer_id", "total_revenue", "order_count"}

Testing UDFs

User-Defined Functions are common sources of bugs — they often don't handle nulls, edge cases, or encoding issues:

# src/myetl/udfs.py
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType, DoubleType
import re

def _parse_phone(raw: str) -> str | None:
    if raw is None:
        return None
    digits = re.sub(r'\D', '', raw)
    if len(digits) == 10:
        return f"({digits[:3]}) {digits[3:6]}-{digits[6:]}"
    return None

parse_phone_udf = udf(_parse_phone, StringType())

# tests/test_udfs.py
import pytest
from myetl.udfs import _parse_phone, parse_phone_udf
from pyspark.sql import Row

# Test the Python function first — fast, no Spark needed
class TestParsePhonePython:
    def test_formats_10_digit_number(self):
        assert _parse_phone("5551234567") == "(555) 123-4567"
    
    def test_strips_formatting(self):
        assert _parse_phone("(555) 123-4567") == "(555) 123-4567"
        assert _parse_phone("555-123-4567") == "(555) 123-4567"
        assert _parse_phone("555.123.4567") == "(555) 123-4567"
    
    def test_returns_none_for_null(self):
        assert _parse_phone(None) is None
    
    def test_returns_none_for_invalid(self):
        assert _parse_phone("12345") is None  # too short
        assert _parse_phone("1234567890123") is None  # too long

# Then test via Spark (integration check)
def test_parse_phone_udf_handles_nulls_in_dataframe(spark):
    df = spark.createDataFrame([
        Row(raw_phone="5551234567"),
        Row(raw_phone=None),
        Row(raw_phone="invalid"),
    ])
    
    result = df.withColumn("phone", parse_phone_udf(df.raw_phone))
    rows = {r.raw_phone: r.phone for r in result.collect()}
    
    assert rows["5551234567"] == "(555) 123-4567"
    assert rows[None] is None
    assert rows["invalid"] is None

Testing Window Functions

Window functions are notoriously easy to get wrong:

# tests/test_window_functions.py
from myetl.transforms import rank_orders_by_customer

def test_order_ranking_is_per_customer(spark):
    """Rank should restart at 1 for each customer."""
    input_df = spark.createDataFrame([
        Row(order_id=1, customer_id=10, amount=50.0),
        Row(order_id=2, customer_id=10, amount=100.0),  # highest for c10
        Row(order_id=3, customer_id=11, amount=200.0),  # only order for c11
        Row(order_id=4, customer_id=10, amount=75.0),
    ])
    
    result = rank_orders_by_customer(input_df)
    rows = {r.order_id: r.rank for r in result.collect()}
    
    # Customer 10's highest-value order gets rank 1
    assert rows[2] == 1  # amount=100
    assert rows[4] == 2  # amount=75
    assert rows[1] == 3  # amount=50
    # Customer 11 only has one order — should be rank 1, not rank 4
    assert rows[3] == 1

Performance: Keeping Tests Fast

# conftest.py — global optimizations
@pytest.fixture(scope="session")
def spark():
    spark = (SparkSession.builder
             .master("local[1]")       # single-threaded for tests
             .config("spark.sql.shuffle.partitions", "1")
             .config("spark.executor.memory", "512m")
             .config("spark.driver.memory", "1g")
             .config("spark.ui.enabled", "false")  # no web UI
             .getOrCreate())
    yield spark
    spark.stop()

For test data, keep it minimal:

# Bad: read 10GB of production data
df = spark.read.parquet("s3://prod/orders/year=2026/")

# Good: 10-20 rows covering the cases you need
input_df = spark.createDataFrame([
    Row(order_id=1, ...),
    Row(order_id=2, ...),
])

A unit test suite for PySpark should run in under 2 minutes. If it's slower, you're reading too much data.

CI Integration

# .github/workflows/pyspark-tests.yml
name: PySpark Tests
on: [push, pull_request]

jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: '3.11'
      
      - name: Install Java (required for Spark)
        uses: actions/setup-java@v4
        with:
          java-version: '11'
          distribution: 'temurin'
      
      - name: Install dependencies
        run: |
          pip install pyspark==3.5.1 pytest chispa pytest-cov
      
      - name: Run PySpark tests
        run: |
          pytest tests/ -v \
            --cov=src/myetl \
            --cov-report=xml \
            --timeout=120
        env:
          PYSPARK_PYTHON: python3
          JAVA_HOME: ${{ env.JAVA_HOME }}

The key discipline: every transformation function should be independently testable with spark.createDataFrame(rows) as input. If you can't do that, the function is too tightly coupled to the infrastructure and needs to be refactored before it can be tested.

PySpark Testing with pytest: Unit Tests for DataFrame Transformations

HelpMeTest

The Core Problem: Coupling to SparkSession

Setting Up pytest with PySpark

Testing DataFrame Transformations

Schema Validation Tests

Testing UDFs

Testing Window Functions

Performance: Keeping Tests Fast

CI Integration

Read more

Test Result Reporting and Failing Fast in CI Pipelines

SOC 2 Evidence Collection Automation: Stop Emailing Screenshots to Your Auditor

PCI-DSS 4.0 Penetration Testing: Scoping, Methodology, and What Assessors Actually Check

Data Pipeline CI/CD Testing: Automate Quality Gates