Skip to content
Cannot retrieve contributors at this time
171 lines (140 sloc) 5.37 KB
from unittest import TestCase
from ipykernel.tests import utils
from nbformat.converter import convert
from nbformat.reader import reads
import re
import json
from copy import copy
import unittest
# py3
from queue import Empty
def isstr(s):
return isinstance(s, str)
except ImportError:
# py2
from Queue import Empty
def isstr(s):
return isinstance(s, basestring) # noqa
class NotebookTest(TestCase):
"""Validate a notebook. All code cells are executed in order. The output is either checked
for errors (if no reference output is present), or is compared against expected output.
Useful references:
IGNORE_TYPES = ["execute_request", "execute_input", "status", "pyin"]
STRIP_KEYS = ["execution_count", "traceback", "prompt_number", "source"]
def _test_notebook(self, notebook, test):
with open(notebook) as f:
nb = convert(reads(, self.NBFORMAT_VERSION)
_, kernel = utils.start_new_kernel()
for i, c in enumerate([c for c in nb.cells if c.cell_type == 'code']):
self._test_notebook_cell(self.sanitize_cell(c), i, kernel, test)
def _test_notebook_cell(self, cell, i, kernel, test):
if hasattr(cell, 'source'): # nbformat 4.0 and later
code = cell.source
code = cell.input
iopub = kernel.iopub_channel
outputs = []
msg = None
no_error = True
first_error = -1
error_msg = ''
while self.should_continue(msg):
msg = iopub.get_msg(block=True, timeout=1)
except Empty:
if msg['msg_type'] not in self.IGNORE_TYPES:
if msg['msg_type'] == 'error':
error_msg = ' ' + msg['content']['ename'] + '\n ' + msg['content']['evalue']
no_error = False
if first_error == -1:
first_error = i
i = len(outputs)
expected = i < len(cell.outputs) and cell.outputs[i] or []
o = self.transform_message(msg, expected)
if (test == 'check_error'):
self.assertTrue(no_error, 'Executing cell %d resulted in an error:\n%s'%(first_error, error_msg))
# Compare computed output against stored output.
# TODO: This doesn't work right now as the generated output is too diverse to
# be verifiable.
scrub = lambda x: self.dump_canonical(list(self.scrub_outputs(x)))
scrubbed = scrub(outputs)
expected = scrub(cell.outputs)
#self.assertEqual(scrubbed, expected, "\n{}\n\n{}".format(scrubbed, expected))
def dump_canonical(self, obj):
return json.dumps(obj, indent=2, sort_keys=True)
def scrub_outputs(self, outputs):
remove all scrubs from output data and text
for output in outputs:
out = copy(output)
for scrub, sub in []:#self.scrubs.items():
def _scrubLines(lines):
if isstr(lines):
return re.sub(scrub, sub, lines)
return [re.sub(scrub, sub, line) for line in lines]
if "text" in out:
out["text"] = _scrubLines(out["text"])
if "data" in out:
if isinstance(out["data"], dict):
for mime, data in out["data"].items():
out["data"][mime] = _scrubLines(data)
out["data"] = _scrubLines(out["data"])
yield out
def strip_keys(self, d):
remove keys from STRIP_KEYS to ensure comparability
for key in self.STRIP_KEYS:
d.pop(key, None)
return d
def sanitize_cell(self, cell):
remove non-reproducible things
for output in cell.outputs:
return cell
def transform_message(self, msg, expected):
transform a message into something like the notebook
"output_type": {
"pyout": "execute_result",
"pyerr": "error"
output = {
u"output_type": msg["msg_type"]
output = self.strip_keys(output)
for key, swaps in SWAP_KEYS.items():
if key in output and output[key] in swaps:
output[key] = swaps[output[key]]
if "data" in output and "data" not in expected:
output["text"] = output["data"]
del output["data"]
return output
def should_continue(self, msg):
determine whether the current message is the last for this cell
if msg is None:
return True
return not (msg["msg_type"] == "status" and
msg["content"]["execution_state"] == "idle")
You can’t perform that action at this time.