# Testing PySpark Code Locally

> Structure transformations as pure DataFrame functions, run them on a local SparkSession in pytest, and compare results reliably.

- Source: https://akshaythakare.com/notes/pyspark/testing-pyspark-code
- Topic: PySpark
- Level: intermediate
- Reading time: 2 min
- Tags: pyspark, testing, pytest
- Updated: 2026-10-11

## Make transforms testable

Keep reading and writing at the edges. Put logic in functions that take and return DataFrames.

```python
def daily_sales(orders: DataFrame) -> DataFrame:
    return (orders.filter("status = 'paid'")
                  .groupBy(F.to_date("ts").alias("business_date"))
                  .agg(F.sum("amount").alias("net_sales")))
```

## One local SparkSession per test run

```python
import pytest
from pyspark.sql import SparkSession

@pytest.fixture(scope="session")
def spark():
    s = (SparkSession.builder.master("local[2]")
         .config("spark.sql.shuffle.partitions", "2")   # fast for tiny data
         .getOrCreate())
    yield s
    s.stop()
```

Starting Spark takes seconds, so share one session across all tests.

## Build input with an explicit schema

```python
def test_daily_sales_ignores_unpaid(spark):
    rows = [("O1", "paid", 100.0, "2026-01-05 10:00:00"),
            ("O2", "cancelled", 50.0, "2026-01-05 11:00:00")]
    df = spark.createDataFrame(rows, "order_id string, status string, amount double, ts string") \
              .withColumn("ts", F.to_timestamp("ts"))
    result = {r.business_date.isoformat(): r.net_sales for r in daily_sales(df).collect()}
    assert result == {"2026-01-05": 100.0}
```

Collecting a tiny result into a Python dict makes the comparison independent of row order.

> **Tip: What to cover**
>
> Null handling, duplicate keys, empty input and schema mismatches. These break real jobs far more often than the happy path.