Nose plugin to run IPython notebooks as tests (run each notebook and check if any cells throw exceptions)

ipython_nosetest.py

import os
import unittest

import nose
from IPython.kernel import KernelManager
from IPython.nbformat.current import reads
from nose.plugins.base import Plugin
from nose.util import anyp


class KernelProxy(object):
    def __init__(self):
        self.kernel_manager = None
        self.client = None
        self.shell = None

    def __del__(self):
        self.destroy_kernel()

    def setup_kernel(self):
        self.kernel_manager = KernelManager()
        self.kernel_manager.start_kernel(stderr=open(os.devnull, 'w'))
        try:
            self.client = self.kernel_manager.client()
        except AttributeError:
            # 0.13
            self.client = self.kernel_manager
        self.client.start_channels()
        self.shell = self.client.shell_channel

    def destroy_kernel(self):
        try:
            self.client.stop_channels()
            self.kernel_manager.shutdown_kernel()
            del self.kernel_manager
            del self.client
            del self.shell
        except:  # Eat all exceptions when destroying kernel
            pass


class IPynbFileCase(unittest.TestCase):
    _multiprocess_can_split_ = False
    _multiprocess_shared_ = False

    def __init__(self, cell_number, cell, notebook_name, kernel_proxy):
        unittest.TestCase.__init__(self)
        self.cell = cell
        self.kernel_proxy = kernel_proxy
        self.notebook_name = notebook_name
        self.cell_number = cell_number

    def id(self):
        return '{0}[{1}]'.format(self.notebook_name,
                                    self.cell_number)

    def shortDescription(self):
        return self.id()

    def runTest(self):
        self.kernel_proxy.shell.execute(self.cell.input)
        # wait for finish, maximum 20s
        reply = self.kernel_proxy.shell.get_msg(timeout=20)['content']

        if reply['status'] == 'error':
            raise Exception('\n'.join(reply['traceback']))


class IPythonNotebookTester(Plugin):
    """
    Activate IPython Notebook tester
    """
    extension = 'ipynb'
    name = 'ipynb'

    def __init__(self):
        Plugin.__init__(self)
        self.current_kernel = None

    def options(self, parser, env):
        """Register commmandline options.
        """
        Plugin.options(self, parser, env)

    def configure(self, options, config):
        """Configure plugin.
        """
        Plugin.configure(self, options, config)

    def beforeContext(self):
        del self.current_kernel
        self.current_kernel = KernelProxy()
        self.current_kernel.setup_kernel()

    def afterContext(self):
        if self.current_kernel is not None:
            self.current_kernel.destroy_kernel()
            self.current_kernel = None

    def loadTestsFromFile(self, filename):
        """Load ipynb cells and run them as 'tests'

        Tests are loaded only if filename's extension matches 'ipynb'.
        """
        if self.is_notebook(filename):
            basename = os.path.basename(filename)
            with open(filename) as f:
                nb = reads(f.read(), 'json')

            suite = []
            cell_count = 0
            for ws in nb.worksheets:
                for cell in ws.cells:
                    if cell.cell_type != 'code':
                        continue
                    suite.append(IPynbFileCase(cell_count, cell, basename,
                                               self.current_kernel))
                    cell_count += 1
            return suite

    def is_notebook(self, name):
        return anyp(name.endswith, self.extension)

    def wantFile(self, file):
        """Override to select all modules and any file ending with 'ipynb'.
        """
        if self.is_notebook(file):
            return True
        return None


if __name__ == '__main__':
    nose.run('menpo', argv=['--with-ipynb', '-v'],
             addplugins=[IPythonNotebookTester()])