Test Spark code with pytest on Sail
Spark transformation code is ordinary Python functions over DataFrames, so it can be unit tested like any other code. This example tests clean_orders, a function that normalizes emails, keeps the latest version of each order and stores amounts in cents. The tests run against Sail: Java and a cluster are not needed.
Run it
Run all commands below from the repository root (sail-examples/):
uv run migrate/test-spark-code-with-pytest/main.py
Running migrate/test-spark-code-with-pytest/main.py runs the four tests in migrate/test-spark-code-with-pytest/test_orders.py. To run them with pytest directly, install the same packages and run pytest migrate/test-spark-code-with-pytest/test_orders.py from the repository root (sail-examples/).
Output
The newest capture appears first. Older output stays here when a result changes.
Current output · Sail 0.7.3
Captured on Sail 0.7.3. Latest capture: 2026-11-05, Python 3.12.14.
.... [100%]
4 passed in 1.25s
Python code
main.py65 lines
# /// script
# requires-python = ">=3.12"
# dependencies = [
# "pysail==0.7.3",
# "pyspark-client==4.1.3",
# "pyarrow==26.0.0",
# "pandas==3.0.6",
# "pytest==9.1.1",
# ]
# ///
"""Unit-test Spark transformation code with pytest, using Sail as the local engine.
`clean_orders` is ordinary Spark DataFrame code. `test_orders.py` tests it against a Sail
session created once per test run by the fixture in `conftest.py`, with no Java and no cluster.
Running this file runs those tests.
"""
import sys
from pathlib import Path
from pyspark.sql import DataFrame, SparkSession, Window
from pyspark.sql import functions as F
def clean_orders(orders: DataFrame) -> DataFrame:
"""Normalize emails, keep the latest version of each order, and store amounts in cents."""
latest = Window.partitionBy("order_id").orderBy(F.col("updated_at").desc())
return (
orders.withColumn("email", F.lower(F.trim("email")))
.withColumn("version", F.row_number().over(latest))
.filter("version = 1")
.filter(F.col("amount") >= 0)
.select(
"order_id",
"email",
F.round(F.col("amount") * 100).cast("long").alias("amount_cents"),
"updated_at",
)
)
SAMPLE = [
(1, " Ana@Example.com ", 12.5, "2026-01-01 10:00:00"),
(1, "ana@example.com", 13.0, "2026-01-02 09:00:00"),
(2, "BO@example.com", 7.25, "2026-01-01 11:00:00"),
(3, "cy@example.com", -4.0, "2026-01-03 12:00:00"),
(4, "Di@Example.com", 0.1, "2026-01-04 08:30:00"),
]
def sample_orders(spark: SparkSession) -> DataFrame:
return spark.createDataFrame(
SAMPLE, "order_id long, email string, amount double, updated_at string"
).withColumn("updated_at", F.to_timestamp_ntz("updated_at"))
def run(spark: SparkSession) -> DataFrame:
return clean_orders(sample_orders(spark)).orderBy("order_id")
if __name__ == "__main__":
import pytest
here = Path(__file__).parent
sys.exit(pytest.main(["-q", "-p", "no:cacheprovider", str(here / "test_orders.py")]))
conftest.py16 lines
import pytest
from pyspark.sql import SparkSession
@pytest.fixture(scope="session")
def spark():
"""One Sail server and session for the whole test run, started in this process."""
from pysail.spark import SparkConnectServer
server = SparkConnectServer()
server.start()
_, port = server.listening_address
session = SparkSession.builder.remote(f"sc://localhost:{port}").getOrCreate()
yield session
session.stop()
server.stop()
test_orders.py30 lines
from datetime import datetime
from main import clean_orders, sample_orders
def rows(df):
"""Collected rows as sorted tuples: order-insensitive and readable in a failure message."""
return sorted(tuple(row) for row in df.collect())
def test_keeps_the_latest_version_of_each_order(spark):
result = clean_orders(sample_orders(spark)).filter("order_id = 1")
assert [row.amount_cents for row in result.collect()] == [1300]
def test_normalizes_emails(spark):
emails = {row.email for row in clean_orders(sample_orders(spark)).collect()}
assert emails == {"ana@example.com", "bo@example.com", "di@example.com"}
def test_drops_negative_amounts(spark):
assert clean_orders(sample_orders(spark)).filter("order_id = 3").count() == 0
def test_whole_result(spark):
assert rows(clean_orders(sample_orders(spark))) == [
(1, "ana@example.com", 1300, datetime(2026, 1, 2, 9, 0)),
(2, "bo@example.com", 725, datetime(2026, 1, 1, 11, 0)),
(4, "di@example.com", 10, datetime(2026, 1, 4, 8, 30)),
]
How it works
migrate/test-spark-code-with-pytest/conftest.pydefines a session-scopedsparkfixture. It starts a Sail server in the test process withSparkConnectServer()and connects a normal PySpark session to it. Every test in the run shares the one server.migrate/test-spark-code-with-pytest/test_orders.pycallsclean_orderson a small DataFrame. It checks single properties with plain asserts and the whole result by comparing sorted, collected rows.migrate/test-spark-code-with-pytest/main.pyholds the code under test. The same function runs unchanged in production on Spark or Sail; this repository's checks confirm thatclean_ordersreturns the same rows on both engines.
With the pyspark-client package
- With the lightweight
pyspark-clientpackage,pyspark.testing.assertDataFrameEqualfails before it compares anything: it imports pandas-on-Spark and looks for a local Spark installation. Compare collected rows instead, as these tests do.
Use it in your project
Copy migrate/test-spark-code-with-pytest/conftest.py into your test directory and add pysail and a PySpark client to your test dependencies. Tests that take a spark fixture and use the DataFrame and SQL APIs then run on Sail; tests that use spark.sparkContext, RDDs, assertDataFrameEqual with pyspark-client, or JVM-only settings need changes first. Store wall-clock times as TIMESTAMP_NTZ, as here: collect() returns TIMESTAMP values as Python datetimes in the client's local time zone, whatever the session time zone is, so tests that compare them depend on the machine they run on.