top of page

Mastering Unit Testing in Python: An Essential Guide for Data Scientists



Unit testing is an integral part of the software development process, especially in the realm of data science where the accuracy and reliability of results are paramount. In this tutorial, we'll explore the many benefits of unit testing, its applications, and the different types of testing. We'll be using Python, an industry-standard language for data science, along with the popular testing library pytest.


I. The Benefits of Unit Testing


1. Time-saving


Imagine that you're working on a massive jigsaw puzzle. If you regularly check that each piece fits correctly as you go along, you'll avoid the time-consuming mistake of having to disassemble large parts of your puzzle when you discover an error. This analogy applies perfectly to the practice of unit testing in software development.


Unit tests help developers identify and correct bugs early in the development process, saving significant amounts of time that would be otherwise spent debugging and fixing problems later on.

def test_addition():
    assert 2 + 2 == 4

Running this test gives us an output:

.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK


2. Documentation


Let's consider the scenario of you inheriting an old, antique clock from your grandparents. Now, this is not an ordinary clock but an intricate, complex device with multiple gears and cogs. Imagine if it came with a detailed guide showing how each component worked. It would make understanding and maintaining the clock a lot easier, wouldn't it?


That's exactly what unit tests do. They function as an informal form of documentation. When a new developer joins the team or when you return to your code after a long time, unit tests can provide valuable insights into how each function operates, what kind of input it expects, and what output it returns.

def test_row_to_list(input_string, expected_output):
    assert row_to_list(input_string) == expected_output


3. Trust


Just as a certificate of authenticity reassures you about the genuine nature of an artwork, a well-tested piece of software assures users about its reliability. Having unit tests in place allows users and developers to trust the software. They can run these tests and verify whether the functions work as intended.

def test_square_function():
    assert square(5) == 25


4. Reduced Downtime


Think of a security checkpoint at an airport. It screens for any potential threats before they can cause any harm. Similarly, unit tests are like checkpoints for your code, reducing the potential downtime for production systems.


A Continuous Integration (CI) system can run unit tests whenever code is pushed to a repository. If any unit test fails, the CI system can reject the change, thereby preventing any potential harm to the live system.

def test_database_connection():
    assert connect_to_db("test_db") is not None


II. Applications of Unit Testing


We will now consider a real-world example of a linear regression project and the application of unit testing in its development process.


1. Example Project Overview


For our tutorial, we'll be working with a project that involves creating a model to predict housing prices based on various features. This project will be divided into three modules:

  • Data Module: Creates a clean data file from raw data

  • Feature Module: Computes features from the clean data

  • Models Module: Outputs a model for predicting housing prices from the features


2. Unit tests for the Data Module


Let's start by testing our data module. The purpose of the data module is to clean the raw data on housing area and price. We have functions like row_to_list() and convert_to_int() that need to be tested.

def test_row_to_list():
    assert row_to_list("3,4\\\\n") == [3, 4]
    assert row_to_list("\\\\n") is None
    assert row_to_list("3,abc\\\\n") is None

def test_convert_to_int():
    assert convert_to_int("2,081") == 2081
    assert convert_to_int("") is None
    assert convert_to_int("1, 000") is None

Here, we are testing the functions with different kinds of inputs. If they are working as expected, the output of the test run should look something like:

......
----------------------------------------------------------------------
Ran 6 tests in 0.002s

OK

This means all our tests have passed and our functions are ready to be used for data cleaning.


III. Definitions and Types of Testing


Having seen how unit tests are written, it is important to understand how they fit into the broader landscape of software testing.


1. Unit Tests


Imagine building a house. Before you put the walls and roof together, you would make sure that each brick is of the right quality. Unit tests are similar in that they test individual units of code (like functions or methods) in isolation. They ensure that each unit works correctly, catching any issues at an early stage.

def test_addition():
    assert add(2, 3) == 5


2. Integration Tests


Going back to the house building analogy, once you're confident about the quality of individual bricks, you would start building walls, and subsequently, connecting these walls together. This is where integration tests come into play. They test whether different units of your code work well together, ensuring the connections between different pieces of code are functioning as expected.

def test_add_and_multiply():
    assert add_and_multiply(2, 3, 2) == 10


3. End-to-end Tests


Finally, you would want to make sure the house, as a whole, is built correctly. End-to-end tests serve a similar purpose in software testing. They check if the entire software application operates as expected, starting from one end (like an unprocessed data file) and going through all the units until the other end (like the final output model).

def test_data_to_model():
    raw_data = "data.csv"
    model = data_to_model(raw_data)
    assert model is not None
    assert test_model_on_sample_data(model) is True


IV. Focus on Unit Tests


Given the various types of testing, our primary focus in this tutorial will be on unit testing. Why? Unit tests serve as the foundation for good software, much like the bricks of a house. It's these individual components that contribute to the larger structure. So, mastering unit tests is vital to ensure that each function in your codebase is performing as it should.

def test_string_reverse():
    assert string_reverse("abcd") == "dcba"
    assert string_reverse("1234") == "4321"


V. Advanced Unit Testing


In the next part of this tutorial, we will delve deeper into the pytest library, and cover the creation of more advanced unit tests for the features and models modules of our sample project. This will further solidify your understanding and mastery of unit testing in Python.

def test_feature_computation():
    # Future code for testing feature computation

def test_model_prediction():
    # Future code for testing model prediction


V. Advanced Unit Testing: Dive Deeper into PyTest


We've built a solid foundation in understanding unit tests. Let's now dive deeper

and explore how we can use PyTest for more advanced unit testing scenarios.


1. Test Suites and Markers in PyTest


PyTest allows us to group tests together using "markers". This can be extremely helpful when we want to categorize tests and execute only a specific category. Think of it like having different toolboxes for plumbing, woodworking, and electrical work, and only having to open the one you need.

Let's say we have multiple tests for feature extraction and model training and we'd like to run them separately. We could add a custom marker for each of these types of tests as follows:

import pytest

@pytest.mark.features
def test_feature_extraction():
    # Your test code here

@pytest.mark.models
def test_model_training():
    # Your test code here

And to run only the tests marked with "features", we could use the following command in the terminal:

pytest -m features


2. Parameterized Tests in PyTest


Imagine you want to test a function with several different inputs to make sure it behaves as expected in each case. You could write a separate test for each input, but that would quickly become repetitive. PyTest provides a solution for this through "parameterized tests".


Consider a function multiply(x, y) that multiplies two numbers. We want to test this function with several pairs of numbers. Here is how we can achieve that:

import pytest

@pytest.mark.parametrize("x, y, expected", [
    (1, 2, 2),
    (2, 3, 6),
    (3, 5, 15),
    (6, 7, 42)
])
def test_multiply(x, y, expected):
    assert multiply(x, y) == expected

With this single test, PyTest will actually run four separate tests - one for each pair of numbers. If any of these tests fail, it will let us know which pair caused the failure.


3. Handling Exceptions in PyTest


Not all tests are about making sure that a function produces the expected output. Sometimes, we also want to verify that a function fails in the expected way when given incorrect inputs. For instance, if we have a function that only accepts positive numbers, it should raise an exception if given a negative number.

PyTest provides a convenient way to write tests that expect certain exceptions to be raised. Here is how you can do that:

import pytest

def test_divide_by_zero():
    with pytest.raises(ZeroDivisionError):
        divide(1, 0)

This test will pass if and only if divide(1, 0) raises a ZeroDivisionError. If the function doesn't raise an exception, or if it raises a different type of exception, the test will fail.


This concludes our dive into advanced unit testing with PyTest. Armed with these techniques, you're well equipped to handle a wide range of testing scenarios. Remember, good tests are the cornerstone of reliable, maintainable code. As a data scientist, these testing skills will help you ensure that your data pipelines, feature extraction methods, and models are working as expected, giving you confidence in your results.

bottom of page