Pytest pt3 - Writing Cleaner Tests with Parametrization and Temporary Files

Writing Cleaner Tests with Parametrization and Temporary Files

The most common anti-pattern I see in data engineering tests is a single test function that does a series of if/elif checks to cover different scenarios. When one case fails, you get a vague failure report that doesn't tell you which input failed. In the previous posts we built a foundation with fixtures and mocking; now we'll attack two problems that make test suites harder to maintain: repetitive conditional test logic, and test artifacts leaking between runs. By the end of this post, you'll write tests that clearly declare every input-output pair and automatically clean up after themselves.

Parametrize: Kill Conditional Branches in Tests

@pytest.mark.parametrize replaces conditional branching with a declarative table of inputs and expected outputs.

Here's a test for clean_amount that handles nulls, negative values, and formatted strings:

@pytest.mark.parametrize(
    "raw_amount, expected",
    [
        (100.0, 100.0),
        (0.0, 0.0),
        (None, 0.0),
        (-50.0, 0.0),
        ("1,200.50", 1200.50),
        ("invalid", 0.0),
    ],
    ids=["positive", "zero", "none", "negative", "with_comma", "garbage"],
)
def test_clean_amount(raw_amount, expected):
    assert clean_amount(raw_amount) == expected

The ids parameter gives each case a human‑readable name that appears in the test output. If you omit ids, pytest creates names from the parameter values, but those can get long. Explicit ids make failures immediately obvious: instead of test_clean_amount[1,200.50-1200.50] you get test_clean_amount[with_comma].

You can also parametrize over multiple fixtures. Suppose you have a fixture that prepares a DataFrame, and you want to test a transformation with different column names:

@pytest.mark.parametrize(
    "col_name, expected_sum",
    [("amount", 300.0), ("adjusted_amount", 450.0)],
    ids=["amount_col", "adjusted_col"],
)
def test_total(sample_transactions, col_name, expected_sum):
    result = sample_transactions.agg({col_name: "sum"}).collect()[0][0]
    assert result == expected_sum

No more nearly‑identical test functions. One function, one parametrize decorator, clean separation.

Patching Multiple Dependencies

Before we move to temporary files, let's revisit mocking for a moment. When your function calls more than one external dependency, stacking @patch decorators works, but @patch.multiple keeps things tidy. Here's a realistic S3 reader test using patch.multiple:

from unittest.mock import patch, MagicMock

def test_read_from_s3_with_multiple_mocks():
    mock_s3_client = MagicMock()
    mock_s3_client.get_object.return_value = {
        "Body": MagicMock(read=MagicMock(return_value=b"col1,col2\n1,2"))
    }

    with patch.multiple(
        "my_pipeline.reader",
        boto3=MagicMock(client=MagicMock(return_value=mock_s3_client)),
        parse_csv=MagicMock(return_value=[(1, 2)]),
    ):
        from my_pipeline.reader import read_data
        result = read_data("bucket", "key")
        assert result == [(1, 2)]

Even cleaner is the decorator approach — two stacked @patch decorators and your test body stays at one indent level:

@patch("my_pipeline.reader.parse_csv", return_value=[(1, 2)])
@patch("my_pipeline.reader.boto3.client")
def test_read_data(mock_boto_client, mock_parse_csv):
    mock_s3 = mock_boto_client.return_value
    mock_s3.get_object.return_value = {"Body": MagicMock(read=lambda: b"col1,col2\n1,2")}
    result = read_data("bucket", "key")
    assert result == [(1, 2)]

The order of decorators is bottom‑up: the parameter closest to the function corresponds to the outermost decorator. It reads well once you get used to it.

Dealing with Temporary Files Without Leaving a Mess

Data pipelines often need to write intermediate files — Parquet datasets, CSVs, or checkpoints. Tests that exercise these writes should never leave behind orphaned files that pollute your working directory or, worse, collide with the next test run. Pytest provides built‑in fixtures that create temporary directories, automatically cleaning them up when the test finishes.

The tmp_path Fixture

The most useful tool is tmp_path, a function‑scoped fixture that returns a pathlib.Path pointing to a unique, empty directory. Use it whenever your test needs to write a file:

def test_write_parquet(spark, tmp_path):
    df = spark.range(10)
    output_dir = tmp_path / "output.parquet"
    df.write.parquet(str(output_dir))

    # Verify the data was written correctly
    read_back = spark.read.parquet(str(output_dir))
    assert read_back.count() == 10

After the test completes, pytest recursively deletes the entire temporary directory — no need to hunt down leftover files in /tmp.

tmp_path is function‑scoped by default, which means each test gets its own pristine directory. This isolation prevents one test's output from interfering with another's.

Session‑ and Module‑Scoped Temporary Directories

Sometimes you need a directory that persists across multiple tests — for example, a shared location for Spark checkpoints or a collection of reference files used by a whole test module. The tmp_path_factory fixture gives you control over scope and lifetime. It's a session‑scoped fixture that you call inside your own fixture:

@pytest.fixture(scope="module")
def checkpoint_dir(tmp_path_factory):
    path = tmp_path_factory.mktemp("checkpoints")
    yield path
    # Pytest will clean up the entire factory directory at the end of the session,
    # but you can also explicitly remove it here if you need to reclaim space early.
    import shutil
    shutil.rmtree(path, ignore_errors=True)

The mktemp() method creates a named temporary directory that lives until the session ends. Because it's module‑scoped, the same directory is reused for all tests in that module, and then it's wiped after the last test finishes. If your tests produce very large files, you might want to use a fixture like this with explicit cleanup to avoid running out of disk space during the session.

Avoiding Hardcoded Paths

Resist the temptation to use hardcoded paths like /tmp/my‑test‑data. They survive test runs, cause collisions when multiple developers or CI jobs share a machine, and can lead to frustrating "it worked on my machine" bugs. Always use tmp_path or tmp_path_factory — they make your tests self‑contained and deterministic.

Conclusion

Parametrization and temporary file handling solve two of the biggest paper cuts in data engineering tests: repetitive conditional logic that hides which case failed, and orphaned files that pollute your environment. With @pytest.mark.parametrize, every input-output pair becomes a named, independently-reported test case. With tmp_path and tmp_path_factory, every file your test creates vanishes when the test ends.

In the final post of this series, we'll look at how to run these tests effectively — colorful logs, splitting unit and integration tests with markers, and auto‑skipping rules that keep CI green without sacrificing local coverage.