Skip to content

Testing

You can unit test your transformations in the same way you would normally but now because you have already defined the input and output schemas for your production code you can now use them as a convenient way to produce input fixtures for your tests:

from dataclasses import dataclass
from typing import Annotated

import polars as pl
from polars.testing import assert_frame_equal

from data_joinery import Schema, Strict, transform


@dataclass
class Customer:
    customer_id: str
    name: str
    is_active: bool


@transform
def filter_active_customers(
    customers: Annotated[pl.LazyFrame, Strict(Customer)],
) -> Annotated[pl.LazyFrame, Strict(Customer)]:
    return customers.filter(pl.col("is_active"))


def test_filter_active_customers():
    customer_schema = Schema(Customer)

    customers = customer_schema.create_dataframe(
        [
            Customer("1", "Alice", True),
            Customer("2", "Bob", False),
            Customer("3", "Charlie", True),
        ],
        pl.LazyFrame,
    )
    expected_active_customers = customer_schema.create_dataframe(
        [
            Customer("1", "Alice", True),
            Customer("3", "Charlie", True),
        ],
        pl.LazyFrame,
    )

    active_customers = filter_active_customers(customers)
    assert_frame_equal(active_customers.collect(), expected_active_customers.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/transform/testing/test_example_polars.py .                [100%]

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

from dataclasses import dataclass
from typing import Annotated

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

from data_joinery import Schema, Strict, transform


@dataclass
class Customer:
    customer_id: str
    name: str
    is_active: bool


@transform
def filter_active_customers(
    customers: Annotated[DataFrame, Strict(Customer)],
) -> Annotated[DataFrame, Strict(Customer)]:
    return customers.filter(customers.is_active)


def test_filter_active_customers():
    spark = SparkSession.builder.getOrCreate()
    customer_schema = Schema(Customer)

    customers = customer_schema.create_dataframe(
        [
            Customer("1", "Alice", True),
            Customer("2", "Bob", False),
            Customer("3", "Charlie", True),
        ],
        DataFrame,
        session=spark,
    )
    expected_active_customers = customer_schema.create_dataframe(
        [
            Customer("1", "Alice", True),
            Customer("3", "Charlie", True),
        ],
        DataFrame,
        session=spark,
    )

    active_customers = customers.transform(filter_active_customers)
    assertDataFrameEqual(active_customers, expected_active_customers)

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/transform/testing/test_example_spark.py .                 [100%]

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