Small Data, Real Spark: A Fast Local Unit-Testing Loop with Docker

01. Why Mocking Spark Isn't Enough

Spark's API surface is large, and its behavior is genuinely surprising in places: null handling in SQL expressions, implicit type coercion, how a function's closure gets serialized to run somewhere else. A mock can only enforce what you already believe to be true, it can't catch you being wrong about how Spark itself behaves.

• Column and SQL semantics: null propagation, type coercion, and edge cases are easy to get wrong from memory, and invisible behind a mock that just returns whatever you told it to.
• Serialization: closures and functions shipped to executors behave differently once a real process boundary is involved, a mock never crosses one.
• API drift: hand-maintained mocks quietly fall out of sync with the real API's behavior across Spark version upgrades.

None of this requires a large cluster or realistic data volumes to catch. It requires a real Spark engine, and a handful of rows chosen deliberately to hit the specific case under test.

02. Separate the Computation from the I/O

The pattern only works if there's something worth testing with a real, tiny Spark session in the first place. That means keeping the actual transformation logic, filtering, joining, aggregating, deriving columns, in functions with a simple shape: take one or more DataFrames in, return a DataFrame out. No file paths, no buckets, no database connections inside them.

def normalize_prices(df):
    return df.withColumn("price", df.price.cast("double"))

That function gets tested with a real, tiny SparkSession and a couple of hand-picked rows, one that should convert cleanly, one that shouldn't:

from transform import normalize_prices
from spark_test_util import build_spark
def test_normalize_prices():
    spark = build_spark()
    df = spark.createDataFrame(
        [("A", "9.99"), ("B", "n/a")],
        ["sku", "price"],
    )
    result = normalize_prices(df).collect()
    assert result[0].price == 9.99
    assert result[1].price is None

Everything that touches the outside world, loading data from a path, writing a result back out, calling another service, belongs in a separate layer: a wrapper function or a class that orchestrates calls to the computation functions and to the I/O. That layer has almost no logic of its own, it just wires things together, so it doesn't need a real Spark session to test. Plain mocks are enough, and appropriate, to assert that this wrapper called read with the right path and write with the right destination, without ever touching a real DataFrame:

class PriceImportPipeline:
    def __init__(self, spark, storage):
        self.spark = spark
        self.storage = storage
    def run(self, input_path, output_path):
        df = self.storage.read(self.spark, input_path)
        result = normalize_prices(df)
        self.storage.write(result, output_path)

from unittest.mock import Mock
from pipeline import PriceImportPipeline
def test_run_reads_input_and_writes_result():
    spark = Mock()
    storage = Mock()
    pipeline = PriceImportPipeline(spark, storage)
    pipeline.run("input.csv", "output.parquet")
    storage.read.assert_called_once_with(spark, "input.csv")
    storage.write.assert_called_once()

This split is what makes both kinds of tests fast and honest. The computation functions get tested against a real Spark engine with a handful of purpose-built rows, because that's where getting Spark's actual behavior right matters. The orchestration layer gets tested with mocks, because there's no computation left in it to get wrong, only wiring.

03. Building a Disposable Cluster with Docker Compose

A Spark cluster for this kind of testing doesn't need a custom build or a lot of resources. Three services from the official Spark image are enough:

• A master service, exposing the standard cluster and UI ports.
• One or two worker services with modest, explicit CPU and memory limits, the goal is a real engine, not production-scale capacity.
• A test-runner service, built from a small custom image, that mounts the working directory and runs the test suite against tiny fixture files as its command.

services:
  spark-master:
    image: apache/spark:3.5.5
    environment:
      - SPARK_MODE=master
    ports:
      - "7077:7077"
  spark-worker:
    image: apache/spark:3.5.5
    depends_on: [spark-master]
    environment:
      - SPARK_MODE=worker
      - SPARK_MASTER_URL=spark://spark-master:7077
      - SPARK_WORKER_CORES=1
      - SPARK_WORKER_MEMORY=1G
  test-runner:
    build:
      context: .
      dockerfile: docker/Dockerfile
    depends_on: [spark-master, spark-worker]
    environment:
      - SPARK_MASTER_URL=local[*]
      - IN_DOCKER=1
      - PYTHONPATH=/work/src
    volumes:
      - ./:/work
    working_dir: /work
    command: ["pytest", "-q"]

The test-runner's own image just needs Python, a JVM for Spark to run on, and the project's dependencies:

FROM python:3.11-slim
RUN apt-get update && apt-get install -y --no-install-recommends \
    openjdk-17-jre-headless \
  && rm -rf /var/lib/apt/lists/*
ENV JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64
ENV PATH="$PATH:$JAVA_HOME/bin"
COPY docker/requirements.txt /tmp/requirements.txt
RUN pip install --no-cache-dir -r /tmp/requirements.txt
WORKDIR /work

Because the whole cluster boots from a couple of lightweight containers, the identical setup runs unchanged in the CI pipeline as on a laptop: same command, same containers, no separate "CI-only" harness to maintain. And because it's disposable and fast, it fits a genuinely tight coding loop: change the function, rerun the suite against a few sample rows, see results in seconds, and tear the cluster down when done.

Keeping an interactive variant of the test-runner service around, one that drops into a shell instead of firing off pytest immediately, pays for itself the first time a test fails for a reason that isn't obvious from the log alone.

04. Writing Test Code That Runs Anywhere

A small helper function builds the SparkSession: it reads the master URL from the environment, falling back to a local single-process session when nothing else is set. Every test just asks for a session and gets whichever one is correct for where it's currently running:

import os
from pyspark.sql import SparkSession
def build_spark(app="tests"):
    master = os.getenv("SPARK_MASTER_URL", "local[*]")
    return (
        SparkSession.builder
        .appName(app)
        .master(master)
        .getOrCreate()
    )

Bake connectors in at build time

Any storage connectors and their jars belong in the test-runner image itself, fetched once at build time, not downloaded fresh on every run. A flaky download partway through a test suite looks exactly like a real bug until someone notices it's a network blip.

Tune for feedback speed

The cluster's defaults should optimize for a fast loop: fewer shuffle partitions, short network timeouts, the Spark UI switched off. This isn't how a job should be tuned in production, it's how the test cluster should be tuned so a broken test fails in seconds.

05. Pitfalls and Best Practices

Most of the pain in this pattern comes from a handful of repeatable mistakes, not from Spark itself.

• Keep sample data small and specific: a fixture with exactly the two or three rows needed to prove one edge case beats a "realistic" chunk of real data, it's faster, and it doesn't quietly smuggle production data into the test suite. Validating against real, full-sized data belongs to a later, separate stage, not this loop.
• Keep computation functions free of I/O: the moment a path or a client sneaks into a function you want to unit test with real Spark, that test stops being small and fast, and starts needing a real environment underneath it.
• Quarantine flaky tests visibly: a test that fails intermittently and gets silently skipped or deleted takes its coverage with it. Marking it and leaving a clear note about why keeps the gap visible until someone has time to fix it properly.
• Watch for hidden ordering dependencies: a suite where one module's output feeds another's input works by accident until the run order changes. Either remove the dependency or make it explicit.
• Pin versions in the image: a test cluster that silently picks up a new base image or dependency version stops being reproducible.
• Keep scratch tests out of the shared suite: a test written to check one person's one-off scenario is useful in the moment and a liability the moment everyone forgets why it's there.

A visible quarantine beats a silent skip:

@pytest.mark.skip(reason="Flaky: intermittent rounding mismatch, see TICKET-123")
def test_currency_conversion():
    ...

None of these are Spark-specific. They're the same discipline any team applies to a shared test suite, it's just that a real cluster makes the cost of skipping that discipline more expensive, not less.