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

To unit test PySpark code, move each transformation into a function that takes a DataFrame and returns one, with no file paths or table names inside. Start 1 local SparkSession for all the tests with a pytest fixture, build a few rows of input in each test with createDataFrame, and compare the result with pyspark.testing.assertDataFrameEqual, which ignores row order by default and reports the rows that differ.

- Page: https://tecyfy.com/data-engineering/course/testing-and-ci-cd
- Course: Data engineering course (https://tecyfy.com/data-engineering/course), Module 11: Testing, CI/CD and bundles, Unit tests for PySpark with pytest
- Parts of this module: Unit tests for PySpark with pytest (https://tecyfy.com/data-engineering/course/testing-and-ci-cd); Data contracts, a check before the load (https://tecyfy.com/data-engineering/course/testing-and-ci-cd/data-contracts); Data quality over time (https://tecyfy.com/data-engineering/course/testing-and-ci-cd/data-quality); Declarative Automation Bundles, dev and prod (https://tecyfy.com/data-engineering/course/testing-and-ci-cd/bundles); Build it: tests and deploys in CI/CD (https://tecyfy.com/data-engineering/course/testing-and-ci-cd/build-it)

At the end of this part you can:

- Split a PySpark job into functions that can be tested without files, tables or a cluster
- Write pytest tests for them with a shared SparkSession fixture and assertDataFrameEqual
- Choose what a unit test should cover, and what belongs to the checks of later parts

[Module 10](/data-engineering/course/orchestration) 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`](/data-engineering/course/delta-lake-and-iceberg). 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`:

```python
"""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:

```python
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`:

```python
"""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`:

```python
"""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`:

```bash
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.

## Check yourself

1. A test builds its expected DataFrame with the rows in a different order from the function's result. With default settings, what does assertDataFrameEqual do?

   A. Fails, because the rows are in a different order
   B. Passes if the rows are the same, because checkRowOrder is False by default
   C. Sorts both by the first column and fails if they differ
   D. Raises an error asking for checkRowOrder

2. Which of these is a unit test of the shop's load, rather than a check of the data?

   A. Every order in the till's file for 15 September is in orders_daily
   B. orders_daily has a row for every day of September
   C. The amount column is never negative in production
   D. Given 2 rows with the same order_id, orders_for_day keeps 1

### Answers

1. B. assertDataFrameEqual compares the rows as a set unless checkRowOrder=True, which is right for most Spark results, since Spark does not promise an order without an orderBy. Set checkRowOrder=True only when the order is part of what the function promises.
2. D. A unit test runs the code on input it made up and checks what the code does with it, here a resent order. The other 3 are about real data on a real day, which is what the checks of parts 2 and 3 watch, every run.

## Questions

### How do I unit test PySpark code?

Put each transformation in a function that takes and returns a DataFrame. In a test file, start a local SparkSession once with a pytest fixture of scope session, build the input with spark.createDataFrame and a few rows, call the function and compare its result with pyspark.testing.assertDataFrameEqual. Run the tests with pytest, on your computer and in CI.

### What is assertDataFrameEqual in PySpark?

A function in pyspark.testing, since Spark 3.5, that checks 2 DataFrames have the same schema and rows. It ignores row order unless checkRowOrder=True, compares floats within rtol and atol, and on a difference raises an error that shows the rows that differ. assertSchemaEqual compares only the schemas.

## Sources

- [Testing PySpark](https://spark.apache.org/docs/latest/api/python/getting_started/testing_pyspark.html)
- [pyspark.testing.assertDataFrameEqual](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.testing.assertDataFrameEqual.html)
- [How to use fixtures](https://docs.pytest.org/en/stable/how-to/fixtures.html)
- [Unit testing for Databricks notebooks](https://docs.databricks.com/aws/en/notebooks/testing)