Getting Started Unit Testing with Pytest | by Ben Bogart | Towards Data Science

This week, I worked on a data ingestion pipeline that has to deal with many different data scenarios. We’ll just say it's government data, and you can imagine the rest. I ran into a problem you may be familiar with. I started updating my functions to handle different data cleaning scenarios, which broke previously working code. In my case, these issues mostly came up using regular expressions to parse text.
At this point, I realized that it was time for unit tests. Why? Because unit tests are designed for exactly this problem. You write small isolated tests to make sure your code is performing as expected. You can easily run these tests whenever you make changes to your code to make sure you haven’t broken anything, and you can continue to add tests to make sure new changes to your code don't break scenarios you have long since forgotten about.
I don’t think unit tests are necessary for every Data Science project. Many projects are small and targeted enough that tests aren’t necessary. But once you are starting to write a package that will require maintenance, which is any package used in production, tests become a huge time saver. It's worth noting at this point that nearly all of the most common data science packages (includingPandas, NumPy, scikit-learn, and matplotlib use pytest as their testing framework.
Python comes packaged with unittest, which works fine in most cases; however, most developers today are using pytest. It's worth taking a look at the documentation for both, but because this is intended to be a crash course, we will only look at pytest and three of its most common tools:
Basic assertion testingParameterizationFixtures
Basic assertion testing
Parameterization
Fixtures
All of the code used here is available in this GitHub repository.
Let’s dig in.
To follow along here, you will need to have pytest installed. On most systems, you can do this by running the following code at the command line:
pip install pytest
We will start at the beginning. Our very first test will test a simple addition function. We will put this code in a file called project.py.
def add(a, b):    return a + b
The test code goes in a file prefixed with test_. When we run our tests, pytest will perform a discovery process. The discovery process will scan the current folder and all of its subfolders recursively for any file that starts with test_ or ends with _test. Pytest will then run the tests located in those files.
Now we know that we need a new file prefixed with test_. We will add the name of the file we are writing the test for after the prefix. That gives us a file name of test_project.py. Add the following to your new test file.
def test_add():    import project    # test basic functionality    assert project.add(1,1) == 2    # try a border case    assert project.add(-1,-1) == -2
Notice that the test function is also prefixed with test_. When pytest scans our file, it will run any functions prefixed with test_.
Looking at the test function, we can see that the actual test employs the assert keyword. If the expression following assert is true, the function will continue. If the expression following assert is false assert will raise an AssertionError. Which will cause the test to fail.
To run the tests, we go to the command line and type pytest.
What would happen if our test failed? Let's find out by changing our first assertion to:
def test_add():    import project    # 1+1 != 3. This will rais an AssertionError    assert project.add(1,1) == 3
Not only does pytest make it clear that our test failed, but it also tells us exactly where by showing us the function and how the test failed! Inspecting this result starting at the bottom, we see:
FAILED test_project.py::test_add — AssertionError: assert 2 == 3 which tells us the name of the file, the test function, and what caused the failure. In this case, our expression, which evaluated to 2 == 3, was false.
Just above that, there is a “test_add” section that finishes with “test_project.py:10: AssertionError”. This tells us the file and the line number of the error.
Everything above that gives us the context of the error so we can understand and fix any problems we encounter quickly.
Notice that although there were two assert statements, pytest considers this only one test because they ran in the same function call.
Let's look at another example involving a Pandas DataFrame that might represent something more like we would see in a data science workflow.
Here is the function we want to test. It adds a column to a DataFrame.
def add_col(df, colname, coldata):        # add a column to a DataFrame    df[colname] = coldata    return df
To test this, we will:
Create another function prefaced with test_ in our test_project.py file.
Define a starting DataFrame.
Define our expected result DataFrame.
Compare the result of adding a column to our starting DataFrame using the add_col function with our expected result DataFrame by calling the DataFrame’s equals method.
def test_add_col():    import project    # start    df = pd.DataFrame(                    [[0, 1], [2, 3]],                    index=['cat', 'dog'],                    columns=['weight', 'height'])    # expected    df_expected = pd.DataFrame(                       [[0, 1, 2], [2, 3, 2]],                       index=['cat', 'dog'],                       columns=['weight', 'height', 'n_ears'])    # test the result    assert project.add_col(df, 'n_ears', 2).equals(df_expected)
Rerunning pytest, we can see that our new test passes, but the previous test still fails. We can see how many of our tests failed and how many passed. The testing doesn’t stop when one of the tests fails; it runs all of our isolated tests and reports back to us on each of them.
Let's get that pretty green line back by fixing this line in our test_add function.
assert project.add(1,1) == 2
Yes!
Another tool that is useful in some cases is the ability to call the test function with different parameters. We will create a new function to test, a subtraction function.
def subtract(a, b):    return a - b
In pytest, parameterization is done with a decorator that defines the parameters we want to pass to our test function. Let’s define a test for our subtract function that uses different parameters.
# test parameterization@pytest.mark.parametrize("a,b,expected", [[2,1,1], [-1,1,-2]])def test_subtract(a, b, expected):     assert project.subtract(a,b) == expected
Let's dissect that. The decorator tells pytest to pass in the elements of the lists as a, band expected to the following function. It will repeat the function call for each list.
If we rerun our test, we will see that there are now 4 tests, not 3. Because parameterization calls the test function multiple times, each function call counts as a separate test.
You can imagine that if you were going to test many functions that perform different tasks on the same DataFrame, it would be convenient to recycle that starting DataFrame between tests. That is what fixtures are for.
A fixture is a function with a special decorator that creates a resource and returns it. We can then call our test function as if we passed in that resource. This is a little counterintuitive, but I think it will make sense once you see it in action.
This is the Fixture.
@pytest.fixture()def df():     df = pd.DataFrame([[0, 1], [2, 3]],                       index=['wallaby', 'kangaroo'],                       columns=['weight', 'height'])return df
When we call our test, we are going to pass in the name of the fixture function. Pytest will call the function and pass its return value to our test function.
def test_add_col2(df):    # expected    df_expected = pd.DataFrame(                       [[0, 1, 3], [2, 3, 5]],                       index=['wallaby', 'kangaroo'],                       columns=['weight', 'height', 'hop_height'])    assert project.add_col(df,                            'hop_height',                           [3,5]).equals(df_expected)
Notice we defined test_add_col2 to take df as a parameter. Because we have defined a fixture function with the name df pytest will call the df fixture function and pass the result into test_add_col2.
Running our tests, we get:
Obviously, there is much more you can do with pytest unit testing, but it's worth getting started now before you know all that stuff. As you see the power of having a suite of tests that you can quickly and easily run whenever you change your package, you will also experience huge time savings down the road. Even if you don’t yet see the value, it's likely that if you want to contribute to any open source project, you will need to write tests for your contributions.
With that said, do not feel that you need to go and immediately create full test suites for all of your code; in most cases, that is just not practical. Instead, consider writing tests when you start making changes to your code or when the project's scope gets large enough that it's not always clear what all of the cases are that a function needs to handle.
Go forth, test your code, and spend your new extra time with your family.
Pytest documentation
Test organization from Contributing to Pandas
Unit testing in Python by Aditya Sharma
Continuous Integration with Python: An Introduction
Unit testing on Wikipedia