Unit Testing Data Processing with pytest

1. The Critical Role of Testing in Data Pipelines

Imagine you deploy a small, seemingly harmless change to a data pipeline. A few hours later, you discover that this change introduced a subtle bug that has corrupted a critical dataset used for company-wide reporting. The dashboards are wrong, trust is broken, and now you have the painful task of finding the root cause and backfilling days of data.

This scenario is all too common. As the resources you will see point out, testing is the primary defense against it.

As the article highlights, there are many types of tests (integration, end-to-end, etc.). Our focus today is on unit tests. A unit test verifies the smallest testable piece of your code—typically a single function or method—in complete isolation. By ensuring each individual component works correctly, we build a strong foundation for the reliability of the entire pipeline.

2. Getting Started with pytest

Python has a built-in testing framework called unittest, but the community has widely adopted pytest for its simplicity and powerful features.

Here’s why pytest is a great choice:

• Simple Syntax: Tests are just regular Python functions, and you use the standard assert keyword, making tests clean and easy to read.
• Automatic Test Discovery: pytest automatically finds your test files and functions as long as you follow a simple naming convention.
• Powerful Fixtures: It has a sophisticated system for managing test setup and teardown (e.g., creating sample data or database connections).
• Rich Plugin Ecosystem: It integrates with hundreds of plugins for things like code coverage (pytest-cov) and mocking (pytest-mock).

Setting Up Your First Test

Let's see how simple it is to get started.

1. Installation: In your activated virtual environment, install pytest:

   pip install pytest pandas
   

2. Naming Conventions: pytest looks for files named test_*.py or *_test.py. Inside those files, it runs functions prefixed with test_.

3. Writing a Test: A test is a function that contains one or more assert statements. An assertion checks if a condition is true. If it is, the test passes. If it's false, the test fails and pytest reports an AssertionError.

To understand the basic workflow, the following video provides a clear and concise demonstration.

3. Writing Tests for Data Processing Logic

Now, let's apply these principles to a common data engineering task: cleaning a Pandas DataFrame.

Imagine we have a function that processes customer data. It's supposed to:

1. Ensure the customer_id column exists.
2. Make sure the customer_id is unique and not null.
3. Check that the age column is an integer.
4. Verify that all age values are within a reasonable range (e.g., 18 to 120).

Let's write tests for this logic. The following video focuses specifically on testing data pipelines and provides perfect examples for our scenario.

Let's consolidate what you've just seen into a single example.

Suppose we have this processing function in a file named processing.py:

# processing.py
import pandas as pd

def clean_customer_data(df: pd.DataFrame) -> pd.DataFrame:
    """A simple cleaning function for customer data."""
    # Ensure customer_id is a primary key
    if 'customer_id' not in df.columns:
        raise ValueError("DataFrame must contain a 'customer_id' column.")
    if df['customer_id'].isnull().any():
        raise ValueError("'customer_id' cannot contain nulls.")
    if not df['customer_id'].is_unique:
        raise ValueError("'customer_id' must be unique.")
    
    # Process age
    df['age'] = pd.to_numeric(df['age'], errors='coerce').astype('Int64')
    
    return df

We would create a test file tests/test_processing.py:

# tests/test_processing.py
import pandas as pd
import pytest
from processing import clean_customer_data

def test_column_existence():
    """Tests that the 'customer_id' column exists."""
    data = {'customer_id': [1, 2, 3], 'age': [25, 30, 35]}
    df = pd.DataFrame(data)
    processed_df = clean_customer_data(df)
    assert 'customer_id' in processed_df.columns

def test_customer_id_is_unique_and_not_null():
    """Tests for uniqueness and non-null constraints on customer_id."""
    data = {'customer_id': [1, 2, 3], 'age': [25, 30, 35]}
    df = pd.DataFrame(data)
    # This should run without raising an error
    clean_customer_data(df)

def test_age_is_integer():
    """Tests that the 'age' column is converted to an integer type."""
    data = {'customer_id': [1, 2, 3], 'age': ['25', '30.0', '35']}
    df = pd.DataFrame(data)
    processed_df = clean_customer_data(df)
    assert pd.api.types.is_integer_dtype(processed_df['age'])

# To run these tests, navigate to your project root in the terminal and run:
# pytest

Notice the repetition in creating the pd.DataFrame? Let's fix that.

4. Efficient Testing with Fixtures and Parameterization

Writing the same setup code in every test is inefficient. pytest provides two powerful features to write cleaner, more effective tests: fixtures and parameterization.

Fixtures: Reusable Test Setups

A fixture is a function that pytest runs before your test function. It's perfect for setting up reusable objects like a sample DataFrame, a database connection, or an API client.

Using a fixture, we can refactor our previous tests:

# tests/test_processing.py (refactored with a fixture)
import pandas as pd
import pytest
from processing import clean_customer_data

@pytest.fixture
def sample_customer_df():
    """Provides a sample customer DataFrame for tests."""
    data = {'customer_id': [1, 2, 3], 'age': ['25', '30.0', '35']}
    return pd.DataFrame(data)

def test_column_existence(sample_customer_df): # Fixture is passed as an argument
    processed_df = clean_customer_data(sample_customer_df)
    assert 'customer_id' in processed_df.columns

def test_age_is_integer(sample_customer_df):
    processed_df = clean_customer_data(sample_customer_df)
    assert pd.api.types.is_integer_dtype(processed_df['age'])

This is much cleaner! The data setup is defined once and reused across multiple tests.

Parameterization: Testing Many Cases with One Function

Often, you want to test your function against multiple inputs, especially edge cases. Instead of writing a new test function for each case, you can parameterize a single test function.

Let's apply this to test how our clean_customer_data function handles various bad inputs for the age column.

# tests/test_processing.py (with parameterization)

@pytest.mark.parametrize("age_input, expected_output", [
    ('42', 42),          # Valid string
    (35.0, 35),          # Valid float
    (None, pd.NA),       # Null value
    ('invalid', pd.NA),  # Non-numeric string
])
def test_age_coercion(age_input, expected_output):
    """Tests how the age column is handled with various inputs."""
    data = {'customer_id': [1], 'age': [age_input]}
    df = pd.DataFrame(data)
    processed_df = clean_customer_data(df)
    
    # Use pd.isna for robust NA comparison
    if pd.isna(expected_output):
        assert pd.isna(processed_df['age'].iloc[0])
    else:
        assert processed_df['age'].iloc[0] == expected_output

With this one test function, we've covered four different scenarios!

import pytest
# Assume categorize_spend is in a file named 'transformations.py'
from transformations import categorize_spend

@pytest.mark.parametrize("amount, expected_category", [
    (49.99, 'low'),
    (50, 'medium'),
    (100, 'medium'),
    (100.01, 'high'),
    (0, 'low'),
    (-10, 'low'), # Testing negative values
])
def test_categorize_spend(amount, expected_category):
    assert categorize_spend(amount) == expected_category

5. Isolating Tests with Mocking

A true unit test must be isolated. This means your function's test should not depend on external systems like databases, APIs, or the file system. Why? Because external systems can be slow, unavailable, or have unpredictable states, making your tests unreliable.

Mocking is the technique of replacing these external dependencies with fake, controlled objects during a test.

For example, if your function reads data from a remote CSV file (pd.read_csv("http://example.com/data.csv")), you don't want your test to make a real network request. You would mock the pd.read_csv function to return a predefined DataFrame instead.

Let's say we have a function that reads data from a URL.

# processing.py
import pandas as pd

def read_data_from_url(url: str) -> pd.DataFrame:
    """Reads a CSV from a URL and returns a DataFrame."""
    return pd.read_csv(url)

Here's how we'd test it by mocking pd.read_csv:

# tests/test_processing.py
import pandas as pd
from processing import read_data_from_url

def test_read_data_from_url(mocker): # We need the mocker fixture
    """Tests the URL reading function without making a network request."""
    # Create the fake data we want our mock to return
    fake_csv_data = "col1,col2\nval1,val2"
    mock_df = pd.read_csv(pd.io.common.StringIO(fake_csv_data))

    # Patch 'pandas.read_csv' and set its return value
    mocker.patch('pandas.read_csv', return_value=mock_df)

    # Call our function
    test_url = "http://fake-url.com/data.csv"
    result_df = read_data_from_url(test_url)

    # Assert that our function called pd.read_csv with the right URL
    pd.read_csv.assert_called_once_with(test_url)

    # Assert that the function returned the DataFrame we expect
    pd.testing.assert_frame_equal(result_df, mock_df)

In this test, no network call ever happens. We test our function's logic in complete isolation, making the test fast, reliable, and independent of external services.

Conclusion

You have now completed the Professional Python module! You've gone from Python fundamentals to advanced topics like decorators, profiling, dependency management, and now, automated testing. This final lesson is arguably one of the most important for your career as a data engineer.

Key Takeaways:

• Testing is Essential: Automated tests prevent bugs, ensure data quality, and give you the confidence to refactor and develop pipelines.
• pytest is Powerful and Simple: Use test_*.py files, test_*() functions, and the assert keyword to write tests. Run them from the terminal with the pytest command.
• Fixtures for Setup: Use @pytest.fixture to create reusable test data and resources, keeping your tests clean and DRY (Don't Repeat Yourself).
• Parameterization for Coverage: Use @pytest.mark.parametrize to test multiple scenarios and edge cases with a single test function.
• Mocking for Isolation: Use mocker to replace external dependencies (files, APIs, databases) to create true, fast, and reliable unit tests.

Preview of the Next Module:
With a solid foundation in professional Python development, we are ready to shift our focus from the code that moves data to the structures that hold data. Our next module, Data Modeling for Analytics, is where we start designing the blueprints for data warehouses and analytical databases. We will begin by exploring the fundamental distinction between two major database design philosophies: OLTP (Online Transaction Processing) and OLAP (Online Analytical Processing). Understanding this difference is the first step toward building efficient and scalable data platforms.