Unit testing Python code in Jupyter notebooks

Most of us agree that we should write unit tests, and many of us actually do. This should be especially true for production code, library code, or if you ascribe to test driven development, during the entire development process.

Often Jupyter notebooks with Python are used for data exploration, and so users may not choose (or need) to write unit tests for their notebook code since they typically may be looking at results for each cell as they progress through the notebook, then coming to a conclusion, and moving on. However, in my experience what typically happens with notebooks is soon the code in the notebook moves beyond data exploration and is useful for further work. Or, perhaps the notebook itself produces results that are useful and need to be run on a regular basis. Perhaps the code needs to be maintained and integrated with external data sources. Then it becomes important to ensure that the code in the notebook can be tested and verified. 

In this case, what are our options for unit testing notebook code? In this article I’ll cover several options for unit testing Python code in a Jupyter notebook.

Maybe just don’t do it?

The first option of Jupyter notebook unit testing is to just not do it at all. By this, I don’t mean don’t unit test your code, but rather extract it from the notebook into separate Python modules that you import back into your notebook. That code should be tested the way you usually unit test your code, whether that be with unittest, pytest, doctest, or another unit testing framework. This article won’t cover all those frameworks in detail, but a great choice for python developers is to not test inside their Jupyter notebooks, but to use the rich assortment of testing frameworks already available for Python code, and to move code to external modules as soon as possible in the development process.

OK, so you can test in a notebook

If you end up deciding you want to leave your code inside a Jupyter notebook, there actually are some unit testing options. Before reviewing a few of them, let’s just setup a code example that we might encounter in a Jupyter notebook. Let’s say your notebook pulls some data from an API, calculates some results from it, then produces some graphs and other data summaries that it persists elsewhere. Maybe there’s a function that produces the proper API URL, and we want to unit test that function. This function has some logic that changes the URL format based on the date for the report. Here’s a debugged version.

import datetime
import dateutil

def make_url(date):
    """Return the url for our API call based on date."""

    if isinstance(date, str):
        date = dateutil.parser.parse(date).date()
    elif not isinstance(date, datetime.date):
        raise ValueError("must be a date")
    if date >= datetime.date(2020, 1, 1):
        return f"https://api.example.com/v2/{date.year}/{date.month}/{date.day}"
    else:
        return f"https://api.example.com/v1/{date:%Y-%m-%d}"

Unit testing with unittest

Normally, when we test with unittest we would either put our test methods in a separate test module, or possibly we’d mix those methods inside the main module. Then we’d need to execute the unittest.main method, possibly as the default method inside a __main__ guard. We can basically do the same thing in our Jupyter notebook. We can make a unitest.TestCase class, perform the tests we want, and then just execute the unit tests in any cell. The results of the tests can even be inspected or asserted to include no failures if you want the notebook execution to fail on errors. You just need to save the output of the unittest.main method and inspect it for errors.

import unittest

class TestUrl(unittest.TestCase):
    def test_make_url_v2(self):
        date = datetime.date(2020, 1, 1)
        self.assertEqual(make_url(date), "https://api.example.com/v2/2020/1/1")
        
    def test_make_url_v1(self):
        date = datetime.date(2019, 12, 31)
        self.assertEqual(make_url(date), "https://api.example.com/v1/2019-12-31")

        
res = unittest.main(argv=[''], verbosity=3, exit=False)

# if we want our notebook to stop processing due to failures, we need a cell itself to fail
assert len(res.result.failures) == 0

test_make_url_v1 (__main__.TestUrl) ... ok
test_make_url_v2 (__main__.TestUrl) ... ok

----------------------------------------------------------------------
Ran 2 tests in 0.001s

OK

This turns out to be fairly straightforward, and if you don’t mind comingling code and tests in your notebook, it works fine.

Unit testing with doctest

Another way to include tests in your code is to use doctest. Doctest uses specially formatted code documentation that includes our tests and the expected results. Below is an updated method with this special code documentation included, both for positive and negative test cases. This is a simple way to test and document code in one place, and often will be used in python modules where the main guard will just run the doct test, like this:

if __name__ == __main__:
    doctest.testmod()

Since we’re in a notebook, we will just add this to a cell below where our code is defined, and it will also work. First, here’s our updated make_urlmethod with the doctest comments.

def make_url(date):
    """Return the url for our API call based on date.
    >>> make_url("1/1/2020")
    'https://api.example.com/v2/2020/1/1'
    
    >>> make_url("1-1-x1")
    Traceback (most recent call last):
        ...
    dateutil.parser._parser.ParserError: Unknown string format: 1-1-x1
    
    >>> make_url("1/1/20001")
    Traceback (most recent call last):
        ...
    dateutil.parser._parser.ParserError: year 20001 is out of range: 1/1/20001
    
    >>> make_url(datetime.date(2020,1,1))
    'https://api.example.com/v2/2020/1/1'
    
    >>> make_url(datetime.date(2019,12,31))
    'https://api.example.com/v1/2019-12-31'
    """
    if isinstance(date, str):
        date = dateutil.parser.parse(date).date()
    elif not isinstance(date, datetime.date):
        raise ValueError("must be a date")
    if date >= datetime.date(2020, 1, 1):
        return f"https://api.example.com/v2/{date.year}/{date.month}/{date.day}"
    else:
        return f"https://api.example.com/v1/{date:%Y-%m-%d}"

import doctest
doctest.testmod()

TestResults(failed=0, attempted=5)

Unit testing with testbook

The testbook project is a different take on notebook unit testing. It allows you to refer to your notebooks in pure Python code from outside a notebook. This allows you to use any testing framework you like (for example, pytest, or unittest) in separate Python modules. You may have a situation where allowing users to modify and update notebook code is the best way to keep code updated and to allow for flexibility for end users. But you may prefer that the code still be tested and verified separately. Testbook makes this an option.

First, you have to install it in your environment:

pip install testbook

or in your notebook

%pip install testbook

Now, in a separate python file, you can import your notebook code and test it there. In that file, you’ll create code that looks like the following, and then you’ll use whichever unit testing framework you prefer to actually execute the unit test. You might create the following code in a Python file (say jupyter_unit_tests.py).

import datetime
import testbook

@testbook.testbook('./jupyter_unit_tests.ipynb', execute=True)
def test_make_url(tb):
    func = tb.ref("make_url")
    date = datetime.date(2020, 1, 2)
    assert func(date) == "https://api.example.com/v2/2020/1/1"

In this case, you can now run the tests with any unit testing framework. For example, with pytest, you would just run the following:

pytest jupyter_unit_tests.py

This works as a normal unit test, and the tests should pass. However, in developing this article, I realized that the testbook code has limited support for passing arguments in the unit test back into the notebook kernel for testing. These arguments are JSON serialized, and the current code knows how to handle a wide array of Python types. But it doesn’t pass a datetime as an object, for example, but as a string. Since our code makes an attempt to parse strings into dates (after I modified it), it works. In other words, the unit test above is not passing in a datetime.date to the make_url method, but rather a string (2020-01-02) that is then parsed into a date. How could you pass in a date from the unit test into the notebook code? You have several options. First, you can make a date object in your notebook just for testing purposes and then refer to that in your unit tests.

testdate1 = datetime.date(2020,1,1)  # for unit test

Then, you could write your unit test to use that variable in the test.

A second option is to inject Python code into the notebook, then refer to it back in your unit test. Both options are shown in the final version of the external unit test. Just save that over jupyter_unit_tests.py and run it using your favorite unit testing framework.

import datetime

import testbook

@testbook.testbook('./jupyter_unit_tests.ipynb', execute=True)
def test_make_url(tb):
    f = tb.ref("make_url")
    d = "2020-01-02"
    assert f(d) == "https://api.example.com/v2/2020/1/2"

    # note that this is actually converted to a string
    d = datetime.date(2020, 1, 2)
    assert f(d) == "https://api.example.com/v2/2020/1/2"

    # this one will be testing the date functionality
    d2 = tb.ref("testdate1")
    assert f(d2) == "https://api.example.com/v2/2020/1/1"

    # this one will inject similar code as above, then use it
    tb.inject("d3 = datetime.date(2020, 2, 3)")
    d3 = tb.ref("d3")
    assert f(d3) == "https://api.example.com/v2/2020/2/3"

Summary

So whether you are a unit testing purist or you just want to sprinkle a few unit tests into your notebooks, there are several options for you to consider. Don’t let your use of notebooks prevent you from doing the right thing in terms of testing your code.