test_notebooks.py

import os
import subprocess
import tempfile

import nbformat

_TEST_DIR = os.path.abspath(os.path.dirname(__file__))
_EXCLUDE = []


def all_notebooks(path="."):
    notebooks = []
    for root, dirs, files in os.walk(path):
        if ".ipynb_checkpoints" in root:
            continue
        for file in files:
            if file.endswith(".ipynb") and (file not in _EXCLUDE):
                notebooks.append(os.path.join(root, file))
    return notebooks


def pytest_generate_tests(metafunc):
    if "notebook" in metafunc.fixturenames:
        metafunc.parametrize("notebook", all_notebooks(_TEST_DIR))


def _notebook_run(path):
    """Execute a notebook via nbconvert and collect output.
       :returns (parsed nb object, execution errors)
    """
    _, notebook = os.path.split(path)
    base, ext = os.path.splitext(notebook)

    with tempfile.NamedTemporaryFile("w", suffix=".ipynb") as fp:
        args = [
            "jupyter",
            "nbconvert",
            "--to",
            "notebook",
            "--execute",
            "--ExecutePreprocessor.kernel_name=python",
            "--ExecutePreprocessor.timeout=None",
            "--output",
            fp.name,
            "--output-dir=.",
            path,
        ]
        subprocess.check_call(args)

        nb = nbformat.read(fp.name, nbformat.current_nbformat, encoding="UTF-8")

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

    return nb, errors


def test_notebook(tmpdir, notebook):
    with tmpdir.as_cwd():
        nb, errors = _notebook_run(notebook)
        assert not errors