Unit tests for functions in a Jupyter notebook?

Solution 1:

Python standard testing tools, such as doctest and unittest, can be used directly in a notebook.

Doctest

A notebook cell with a function and a test case in a docstring:

def add(a, b):
    '''
    This is a test:
    >>> add(2, 2)
    5
    '''
    return a + b

A notebook cell (the last one in the notebook) that runs all test cases in the docstrings:

import doctest
doctest.testmod(verbose=True)

Output:

Trying:
    add(2, 2)
Expecting:
    5
**********************************************************************
File "__main__", line 4, in __main__.add
Failed example:
    add(2, 2)
Expected:
    5
Got:
    4
1 items had no tests:
    __main__
**********************************************************************
1 items had failures:
   1 of   1 in __main__.add
1 tests in 2 items.
0 passed and 1 failed.
***Test Failed*** 1 failures.

Unittest

A notebook cell with a function:

def add(a, b):
    return a + b

A notebook cell (the last one in the notebook) that contains a test case. The last line in the cell runs the test case when the cell is executed:

import unittest

class TestNotebook(unittest.TestCase):
    
    def test_add(self):
        self.assertEqual(add(2, 2), 5)
        

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

Output:

test_add (__main__.TestNotebook) ... FAIL

======================================================================
FAIL: test_add (__main__.TestNotebook)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "<ipython-input-15-4409ad9ffaea>", line 6, in test_add
    self.assertEqual(add(2, 2), 5)
AssertionError: 4 != 5

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

FAILED (failures=1)

Debugging a Failed Test

While debugging a failed test, it is often useful to halt the test case execution at some point and run a debugger. For this, insert the following code just before the line at which you want the execution to halt:

import pdb; pdb.set_trace()

For example:

def add(a, b):
    '''
    This is the test:
    >>> add(2, 2)
    5
    '''
    import pdb; pdb.set_trace()
    return a + b

For this example, the next time you run the doctest, the execution will halt just before the return statement and the Python debugger (pdb) will start. You will get a pdb prompt directly in the notebook, which will allow you to inspect the values of a and b, step over lines, etc.

Note: Starting with Python 3.7, the built-in breakpoint() can be used instead of import pdb; pdb.set_trace().

I created a Jupyter notebook for experimenting with the techniques I have just described. You can try it out with Binder

Solution 2:

I'm the author and maintainer of testbook (a project under nteract). It is a unit testing framework for testing code in Jupyter Notebooks.

testbook addresses all the three approaches that you've mentioned since it allows for testing Jupyter Notebooks as .py files.

Here is an example of a unit test written using testbook

Consider the following code cell in a Jupyter Notebook:

def func(a, b):
    return a + b

You would write a unit test using testbook in a Python file as follows:

import testbook


@testbook.testbook('/path/to/notebook.ipynb', execute=True)
def test_func(tb):
    func = tb.ref("func")

    assert func(1, 2) == 3

Let us know if testbook helps your use case! If not, please feel free to raise an issue on GitHub :)

Features of testbook

Write conventional unit tests for Jupyter Notebooks
Execute all or some specific cells before unit test
Share kernel context across multiple tests (using pytest fixtures)
Inject code into Jupyter notebooks
Works with any unit testing library - unittest, pytest or nose

Links

PyPI GitHub Docs