Skip to content

Add test for Jupyter notebook #230

@clegaspi

Description

@clegaspi

From @montoyjh

import os
import subprocess
import tempfile
import unittest

from glob import glob
from abc import ABCMeta, abstractmethod

import nbformat


class NotebookDirectoryTest(unittest.TestCase, metaclass=ABCMeta):
    """
    This abstract class runs all of the jupyter notebooks in a directory,
    note that a class attribute "directory_name" must be defined

    Example:
        >>> import unittest
        >>> from tri_utils.testing import NotebookDirectoryTest
        >>>
        >>> class MyNotebookTest(NotebookDirectoryTest):
        >>>     directory_name = "my_directory_with_notebooks_in_it"
        >>>
        >>> if __name__ == __main__:
        >>>     unittest.main()
    """

    @property
    @abstractmethod
    def directory_name(self):
        # This is the goofiest way to define an abstract attribute
        # thx mr python
        pass

    def test_all_notebooks_in_directory(self):
        file_list = glob(os.path.join(self.directory_name, "*.ipynb"))
        for filename in file_list:
            # Run jupyter test
            _notebook_run(filename)


def _notebook_run(path):
    """
    Execute a notebook via nbconvert and collect output.
    Taken from
    https://blog.thedataincubator.com/2016/06/testing-jupyter-notebooks/

    Args:
        path (str): file path for the notebook object

    Returns:
        (parsed nb object, execution errors)
    """
    dirname, __ = os.path.split(path)
    os.chdir(dirname)
    with tempfile.NamedTemporaryFile(suffix=".ipynb") as fout:
        args = ["jupyter", "nbconvert", "--to", "notebook", "--execute",
                "--ExecutePreprocessor.timeout=60",
                "--output", fout.name, path]
        subprocess.check_call(args)

        fout.seek(0)
        nb = nbformat.read(fout, nbformat.current_nbformat)

    errors = [output for cell in nb.cells if "outputs" in cell
              for output in cell["outputs"]
              if output.output_type == "error"]

    return nb, errors


if __name__ == "__main__":
    unittest.main()

Executes the notebook and ensures its runs.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions