Unit test PySpark code with pytest: fixtures, assertDataFrameEqual and small test data

Module 10 scheduled the shop's daily load, retried it and backfilled a month. Every one of those runs trusted the code. This module is about earning that trust before the code ships: tests of the code, contracts and checks of the data, and a deployment that runs both. This part needs module 10's load_day.py and module 4's delta_session.py. To rebuild from nothing, run module 10's python make_days.py.

What a unit test is for

A unit test runs 1 piece of code on input you made up and checks what comes out. It needs no real files, no tables and no cluster, so it runs in seconds on your computer and on every change. It answers 1 question: does the code do what I meant, including on the awkward cases? It cannot tell you whether today's file is right; parts 2 and 3 check the data.

In module 10 the load read a file, set its date and wrote a table, all in 1 function. To test the middle step you would need a file and a table. The fix is to split it.

Split the logic from the reading and writing

Make a folder shop with an empty __init__.py, and save this as shop/orders.py:

"""The shop's daily load, split so its logic can be tested without files or tables."""
from datetime import date, timedelta

from pyspark.sql import DataFrame, functions as F

ORDER_SCHEMA = "order_id INT, customer_id STRING, amount DECIMAL(10,2), ordered_at TIMESTAMP"


def day_for_run(run_date: str) -> str:
    """The day a run loads: the run on 2026-09-02 loads 2026-09-01."""
    return str(date.fromisoformat(run_date) - timedelta(days=1))


def orders_for_day(raw: DataFrame, day: str) -> DataFrame:
    """The day's orders as orders_daily stores them: 1 row per order, with its date."""
    return raw.dropDuplicates(["order_id"]).withColumn("order_date", F.lit(day).cast("date"))

load_day.py keeps the reading and the writing and calls these 2 in between:

from shop.orders import ORDER_SCHEMA, day_for_run, orders_for_day

def load_day(spark, run_date: str) -> str:
    day = day_for_run(run_date)
    raw = spark.read.option("header", True).schema(ORDER_SCHEMA).csv(f"landing/orders/{day}.csv")
    orders = orders_for_day(raw, day)
    orders.write.format("delta").mode("overwrite").option("replaceWhere", f"order_date = '{day}'").saveAsTable("orders_daily")
    ...

orders_for_day also drops a repeated order_id, which module 10's load did not. A till that sends an order twice is the kind of case a test is written for, and the next section writes it first.

A Spark session for the tests

Every test needs a Spark session, and starting one takes seconds, so start 1 for the whole run. Save this as tests/conftest.py:

"""Shared by every test in this folder: 1 local Spark session for the whole test run."""
import pytest
from pyspark.sql import SparkSession


@pytest.fixture(scope="session")
def spark():
    session = (
        SparkSession.builder.master("local[1]").appName("tests")
        .config("spark.sql.shuffle.partitions", "1")
        .getOrCreate()
    )
    yield session
    session.stop()

pytest reads conftest.py before the tests in its folder. A test that names spark as an argument gets this session, created the first time a test asks and stopped after the last. local[1] and 1 shuffle partition keep each tiny query fast: the default of 200 partitions is for a cluster.

The tests

Save this as tests/test_orders.py:

"""Unit tests for the shop's daily load."""
from datetime import date
from decimal import Decimal

from pyspark.testing import assertDataFrameEqual, assertSchemaEqual

from shop.orders import ORDER_SCHEMA, day_for_run, orders_for_day


def test_a_run_loads_the_day_before():
    assert day_for_run("2026-09-02") == "2026-09-01"
    assert day_for_run("2026-10-01") == "2026-09-30"  # across a month
    assert day_for_run("2027-01-01") == "2026-12-31"  # across a year


def test_each_order_gets_the_day_it_was_loaded_for(spark):
    raw = spark.createDataFrame([(2001, "C01", Decimal("12.50"), None)], ORDER_SCHEMA)
    expected = spark.createDataFrame(
        [(2001, "C01", Decimal("12.50"), None, date(2026, 9, 1))],
        ORDER_SCHEMA + ", order_date DATE",
    )
    assertDataFrameEqual(orders_for_day(raw, "2026-09-01"), expected)


def test_an_order_sent_twice_is_kept_once(spark):
    raw = spark.createDataFrame(
        [(2001, "C01", Decimal("12.50"), None), (2001, "C01", Decimal("12.50"), None), (2002, "C02", Decimal("8.00"), None)],
        ORDER_SCHEMA,
    )
    result = orders_for_day(raw, "2026-09-01")
    assert result.count() == 2
    assert sorted(row.order_id for row in result.collect()) == [2001, 2002]


def test_the_schema_is_what_orders_daily_expects(spark):
    raw = spark.createDataFrame([], ORDER_SCHEMA)
    expected = spark.createDataFrame([], ORDER_SCHEMA + ", order_date DATE").schema
    assertSchemaEqual(orders_for_day(raw, "2026-09-01").schema, expected)

Each test is a function whose name starts with test_, in a file whose name does too, which is how pytest finds them. Each builds the few rows it needs with createDataFrame and a schema string, so the test shows its whole input on 1 screen. Run them from the folder above shop and tests:

pip install pytest
python -m pytest tests

python -m pytest puts the current folder on Python's path, so from shop.orders import ... works. pytest prints 1 dot for each test that passed and a line with the totals; a failure prints the test, the line and the values. Written before the dropDuplicates, the third test fails and shows that 3 rows came back. That is the order of work: a test that fails for the case you worry about, then the code that makes it pass.

How assertDataFrameEqual compares

assertDataFrameEqual(actual, expected) checks that the 2 DataFrames have the same schema and the same rows. By default it ignores the order of the rows (checkRowOrder=False), because Spark gives no order without orderBy. It compares floating point values within a tolerance (rtol, atol), and it ignores whether a column allows nulls (ignoreNullable=True). On a difference it raises an error that lists the rows that differ, which says far more than a failed assert a.collect() == b.collect(). assertSchemaEqual compares only the schemas, for a test that cares about the columns and types a table expects. Both work in any test framework and run the same on classic Spark and Spark Connect.

What to test, and what not

Test the decisions in your code: the date arithmetic, the rule for a repeated order, a join that must not drop customers, a filter's edges. Keep each test to a handful of rows that you can read. Do not unit test Spark itself, such as whether saveAsTable writes a table, and do not run unit tests on production data. Whether the real file for 15 September arrived, has the right columns and has every order is a question about the data, and it changes every day. Parts 2 and 3 answer it, on every run.