Testing Jupyter Notebook Code with pytest

In this guide, we cover how to test the code inside a Jupyter notebook using pytest. This approach allows you to build comprehensive yet flexible tasks for the user to complete. For example, you can test the contents of a given variable, the return value of a function, or even a class.

There are three primary aspects of testing a Jupyter notebook: the testbook and pytest Python libraries and unit tests. testbook is used to execute the entire notebook and then import "references" to Python objects (variables, classes, and functions) that can be tested using a unit test. pytest provides a simple yet powerful test runner that conforms to typical Python testing approaches, for example, using assert.

To get started, select the Python or Data Science stack for your content. These are the only stacks with testbook and pytest pre-installed. If you need to use a different stack, that's okay! Just add the following as a startup script:

pip3 install --upgrade jupyter_client testbook pytest
apt install -y python-pytest

Now, create a new Jupyter Notebook tab (or use the Jupyter Notebook interface), then a new Python 3 notebook called "Notebook" (the actual filename will be "Notebook.ipynb"). In the notebook, place the following example code:

def double_array(a):
    return [x * 2 for x in a]

This creates a function called double_array that takes an array (a) and returns another array with the values in a doubled.

Once you're done, be sure to click Save to save the tab and notebook file!

Now, create a custom test by selecting the Test a Jupyter Notebook with pytest test template and adding a description. Below is the default code used in the custom test template:

import pytest
from testbook import testbook
​
​
# Set up a shared notebook context to speed up tests.
@pytest.fixture(scope='module')
def tb():
    with testbook('/home/nt-user/workspace/Notebook.ipynb', execute=True) as tb:
        yield tb
​
​
# Test using function call.
def test_double_array(tb):
    double_array = tb.ref("double_array")
    assert double_array([1, 2, 3]) == [2, 4, 6]
​
​
# Test using code injection.
def test_double_array_inject(tb):
    double_array = tb.ref("double_array")
​
    tb.inject("""
        data = [1, 2, 3]
    """)
    data = tb.ref("data")
​
    assert double_array(data) == [2, 4, 6]

Let's take a look at this code. First, it sets up a shared scope for pytest. This will prevents having to re-execute the notebook for each test, which will slow down the test run considerably. Inside of that shared scope on line 8, it executes the notebook in its entirety. You can also specify which cells you want to execute using notebook cell tags (click View > Cell Toolbar > Tags in the notebook to edit cell tags) or a given range of cells. More details can be found here.

Next, it defines two tests. The first one (test_double_array) uses testbook to load a reference to the double_array function from the notebook. This function is then called directly using an array of values. An assertion is made about the return result of that array.

The second test (test_double_array_inject) also loads a reference to the double_array function. It then injects an array directly into the notebook using the tb.inject function. This array is then also referenced and used when calling the function and making the assertion. While this second method does not offer any benefits over the first, it is a good example of how code can easily be injected into a notebook using testbook if necessary.

That's it! All that's left now is to add this check to a task and you are ready to go!