Skip to content

Testing

Pipeline tests check that connected transformations produce the expected result. If a reader, writer, or model cannot run in a test environment, pass a replacement transformation to Pipeline.run() through transform_overrides.

To inspect outputs and writer inputs from a completed run, see Inspecting Pipeline Runs.

Replacing readers and writers

Pass overrides to Pipeline.run() as a mapping from step name to another transformation. Overrides apply only to that run. Give steps you plan to replace explicit names with add_step(..., name="..."). This example replaces the reader and writer, then checks the DataFrame passed to the writer:

from dataclasses import dataclass
from typing import Annotated

import polars as pl
from polars.testing import assert_frame_equal

from data_joinery import Context, Pipeline, Project, transform


@dataclass(frozen=True)
class Storage:
    input_table: str
    output_table: str


@dataclass(frozen=True)
class OrderPipelineContext:
    storage: Storage


@dataclass
class Order:
    order_id: int
    amount: int


@dataclass
class OrderTotal:
    order_id: int
    total: int


@transform
def read_orders(
    storage: Annotated[Storage, Context()],
) -> Annotated[pl.LazyFrame, Project(Order)]:
    return pl.scan_parquet(storage.input_table)


@transform
def calculate_totals(
    orders: Annotated[pl.LazyFrame, Project(Order)],
) -> Annotated[pl.LazyFrame, Project(OrderTotal)]:
    return orders.select("order_id", (pl.col("amount") * 2).alias("total"))


@transform
def write_totals(
    totals: Annotated[pl.LazyFrame, Project(OrderTotal)],
    storage: Annotated[Storage, Context()],
) -> None:
    totals.sink_parquet(storage.output_table)


def build_pipeline() -> Pipeline[OrderPipelineContext]:
    pipeline = Pipeline(OrderPipelineContext)
    read = pipeline.add_step(read_orders, name="read")
    calculate = pipeline.add_step(calculate_totals, name="calculate")
    write = pipeline.add_step(write_totals, name="write")
    read >> calculate >> write
    return pipeline


def test_order_pipeline_end_to_end():
    fixture_orders = pl.LazyFrame({"order_id": [1, 2], "amount": [10, 25]})

    @transform
    def read_fixture() -> Annotated[pl.LazyFrame, Project(Order)]:
        return fixture_orders

    @transform
    def capture_totals(
        totals: Annotated[pl.LazyFrame, Project(OrderTotal)],
    ) -> None:
        pass

    result = build_pipeline().run(
        OrderPipelineContext(Storage("unused", "unused")),
        transform_overrides={
            "read": read_fixture,
            "write": capture_totals,
        },
    )

    written = result.get_one_input("write", pl.LazyFrame)
    expected = pl.LazyFrame({"order_id": [1, 2], "total": [20, 50]})
    assert_frame_equal(written.collect(), expected.collect())

Outputs:

============================= test session starts ==============================
platform linux -- Python 3.12.3, pytest-9.1.1, pluggy-1.6.0
rootdir: /home/runner/work/data-joinery/data-joinery
configfile: pyproject.toml
plugins: platformdirs-4.12.1, layout-enforcer-0.1.0
collected 1 item

docs_src/learn/pipelines/testing/test_pipeline_read_write_overrides_polars.py . [100%]

============================== 1 passed ===============================

from dataclasses import dataclass
from typing import Annotated

import pytest
from pyspark.sql import DataFrame, SparkSession
from pyspark.testing import assertDataFrameEqual

from data_joinery import Context, Pipeline, Project, transform


@dataclass(frozen=True)
class Storage:
    input_table: str
    output_table: str


@dataclass(frozen=True)
class OrderPipelineContext:
    spark: SparkSession
    storage: Storage


@dataclass
class Order:
    order_id: int
    amount: int


@dataclass
class OrderTotal:
    order_id: int
    total: int


@transform
def read_orders(
    spark: Annotated[SparkSession, Context()],
    storage: Annotated[Storage, Context()],
) -> Annotated[DataFrame, Project(Order)]:
    return spark.table(storage.input_table)


@transform
def calculate_totals(
    orders: Annotated[DataFrame, Project(Order)],
) -> Annotated[DataFrame, Project(OrderTotal)]:
    return orders.selectExpr("order_id", "amount * 2 AS total")


@transform
def write_totals(
    totals: Annotated[DataFrame, Project(OrderTotal)],
    storage: Annotated[Storage, Context()],
) -> None:
    totals.write.saveAsTable(storage.output_table)


def build_pipeline() -> Pipeline[OrderPipelineContext]:
    pipeline = Pipeline(OrderPipelineContext)
    read = pipeline.add_step(read_orders, name="read")
    calculate = pipeline.add_step(calculate_totals, name="calculate")
    write = pipeline.add_step(write_totals, name="write")
    read >> calculate >> write
    return pipeline


@pytest.fixture(scope="module")
def spark():
    return SparkSession.builder.master("local[1]").getOrCreate()


def test_order_pipeline_end_to_end(spark: SparkSession):
    fixture_orders = spark.createDataFrame(
        [(1, 10), (2, 25)], "order_id BIGINT, amount BIGINT"
    )

    @transform
    def read_fixture() -> Annotated[DataFrame, Project(Order)]:
        return fixture_orders

    @transform
    def capture_totals(
        totals: Annotated[DataFrame, Project(OrderTotal)],
    ) -> None:
        pass

    result = build_pipeline().run(
        OrderPipelineContext(spark, Storage("unused", "unused")),
        transform_overrides={
            "read": read_fixture,
            "write": capture_totals,
        },
    )

    written = result.get_one_input("write", DataFrame)
    expected = spark.createDataFrame(
        [(1, 20), (2, 50)], "order_id BIGINT, total BIGINT"
    )
    assertDataFrameEqual(written, expected)

Outputs:

============================= test session starts ==============================
platform linux -- Python 3.12.3, pytest-9.1.1, pluggy-1.6.0
rootdir: /home/runner/work/data-joinery/data-joinery
configfile: pyproject.toml
plugins: platformdirs-4.12.1, layout-enforcer-0.1.0
collected 1 item

docs_src/learn/pipelines/testing/test_pipeline_read_write_overrides_spark.py . [100%]

============================== 1 passed ==============================

The production reader and writer are not called. An override may remove context dependencies but cannot introduce dependencies absent from the original transformation. The pipeline still receives its complete declared context, keeping its public interface stable. Contract parameters must remain compatible, and the override must declare the same output contract as the production transformation.

Replacing a model

Overrides can also replace a model training step with a small dummy model, so the surrounding steps can be tested:

from dataclasses import dataclass
from typing import Annotated

import polars as pl
from polars.testing import assert_frame_equal

from data_joinery import Pipeline, Project, transform


@dataclass
class Observation:
    value: int


class Model:
    def predict(self, value: int) -> int:
        raise NotImplementedError


class DummyModel(Model):
    def predict(self, value: int) -> int:
        return value * 2


@transform
def read_observations() -> Annotated[pl.LazyFrame, Project(Observation)]:
    return pl.LazyFrame({"value": [1, 2]})


@transform
def train_model(
    observations: Annotated[pl.LazyFrame, Project(Observation)],
) -> Model:
    raise RuntimeError("expensive production training")


@transform
def predict(
    observations: Annotated[pl.LazyFrame, Project(Observation)],
    model: Model,
) -> Annotated[pl.LazyFrame, Project(Observation)]:
    return observations.select(
        pl.col("value").map_elements(model.predict, return_dtype=pl.Int64)
    )


def build_pipeline() -> Pipeline:
    pipeline = Pipeline()
    read = pipeline.add_step(read_observations, name="read")
    train = pipeline.add_step(train_model, name="train")
    prediction = pipeline.add_step(predict, name="predict")
    pipeline.connect(read, train)
    pipeline.connect_many([read, train], prediction)
    return pipeline


def test_model_pipeline_with_dummy_training():
    @transform
    def train_dummy(
        observations: Annotated[pl.LazyFrame, Project(Observation)],
    ) -> Model:
        return DummyModel()

    result = build_pipeline().run(
        transform_overrides={"train": train_dummy},
    )

    predicted = result.get_output("predict", pl.LazyFrame)
    expected = pl.LazyFrame({"value": [2, 4]})
    assert_frame_equal(predicted.collect(), expected.collect())

Outputs:

============================= test session starts ==============================
platform linux -- Python 3.12.3, pytest-9.1.1, pluggy-1.6.0
rootdir: /home/runner/work/data-joinery/data-joinery
configfile: pyproject.toml
plugins: platformdirs-4.12.1, layout-enforcer-0.1.0
collected 1 item

docs_src/learn/pipelines/testing/test_pipeline_model_override_polars.py . [100%]

=============================== warnings summary ===============================
docs_src/learn/pipelines/testing/test_pipeline_model_override_polars.py::test_model_pipeline_with_dummy_training
  /home/runner/work/data-joinery/data-joinery/docs_src/learn/pipelines/testing/test_pipeline_model_override_polars.py:43: PolarsInefficientMapWarning: 
  Expr.map_elements is significantly slower than the native expressions API.
  Only use if you absolutely CANNOT implement your logic otherwise.
  Replace this expression...
    - pl.col("value").map_elements(predict)
  with this one instead:
    + pl.col("value") * 2

    pl.col("value").map_elements(model.predict, return_dtype=pl.Int64)

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
========================= 1 passed, 1 warning =========================

from dataclasses import dataclass
from typing import Annotated

import pytest
from pyspark.sql import DataFrame, SparkSession
from pyspark.testing import assertDataFrameEqual

from data_joinery import Context, Pipeline, Project, transform
from data_joinery.backends.spark import SparkContext


@dataclass
class Observation:
    value: int


class Model:
    def predict(self, value: int) -> int:
        raise NotImplementedError


class DummyModel(Model):
    def predict(self, value: int) -> int:
        return value * 2


@transform
def read_observations(
    spark: Annotated[SparkSession, Context()],
) -> Annotated[DataFrame, Project(Observation)]:
    return spark.createDataFrame([(1,), (2,)], "value BIGINT")


@transform
def train_model(
    observations: Annotated[DataFrame, Project(Observation)],
) -> Model:
    raise RuntimeError("expensive production training")


@transform
def predict(
    spark: Annotated[SparkSession, Context()],
    observations: Annotated[DataFrame, Project(Observation)],
    model: Model,
) -> Annotated[DataFrame, Project(Observation)]:
    predictions = [(model.predict(row.value),) for row in observations.collect()]
    return spark.createDataFrame(
        predictions,
        "value BIGINT",
    )


def build_pipeline() -> Pipeline[SparkContext]:
    pipeline = Pipeline(SparkContext)
    read = pipeline.add_step(read_observations, name="read")
    train = pipeline.add_step(train_model, name="train")
    prediction = pipeline.add_step(predict, name="predict")
    pipeline.connect(read, train)
    pipeline.connect_many([read, train], prediction)
    return pipeline


@pytest.fixture(scope="module")
def spark():
    return SparkSession.builder.master("local[1]").getOrCreate()


def test_model_pipeline_with_dummy_training(spark: SparkSession):
    @transform
    def train_dummy(
        observations: Annotated[DataFrame, Project(Observation)],
    ) -> Model:
        return DummyModel()

    result = build_pipeline().run(
        SparkContext(spark),
        transform_overrides={"train": train_dummy},
    )

    predicted = result.get_output("predict", DataFrame)
    expected = spark.createDataFrame([(2,), (4,)], "value BIGINT")
    assertDataFrameEqual(predicted, expected)

Outputs:

============================= test session starts ==============================
platform linux -- Python 3.12.3, pytest-9.1.1, pluggy-1.6.0
rootdir: /home/runner/work/data-joinery/data-joinery
configfile: pyproject.toml
plugins: platformdirs-4.12.1, layout-enforcer-0.1.0
collected 1 item

docs_src/learn/pipelines/testing/test_pipeline_model_override_spark.py . [100%]

============================== 1 passed ==============================

An override must declare the same output contract as the production transformation. The dummy training transformation therefore declares Model, but its runtime value can still be a DummyModel instance.

Context or an override?

Use Context() when the implementation is suitable for a test and only its values need to change. Paths, dates, thresholds, and other configuration usually belong in the pipeline's context dataclass.

Use an override when the implementation itself should not run, such as:

  • a streaming sink that should be replaced with a memory sink
  • a reader that depends on unavailable external infrastructure
  • an expensive model-training implementation
  • a writer with effects that the test needs to capture