Skip to content
Permalink
 
 
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
try:
# 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:
http://nbformat.readthedocs.org/en/latest/format_description.html
http://jupyter-client.readthedocs.org/en/latest/messaging.html
"""
IGNORE_TYPES = ["execute_request", "execute_input", "status", "pyin"]
STRIP_KEYS = ["execution_count", "traceback", "prompt_number", "source"]
NBFORMAT_VERSION = 4
def _test_notebook(self, notebook, test):
with open(notebook) as f:
nb = convert(reads(f.read()), 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
else:
code = cell.input
iopub = kernel.iopub_channel
kernel.execute(code)
outputs = []
msg = None
no_error = True
first_error = -1
error_msg = ''
while self.should_continue(msg):
try:
msg = iopub.get_msg(block=True, timeout=1)
except Empty:
continue
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)
outputs.append(o)
if (test == 'check_error'):
self.assertTrue(no_error, 'Executing cell %d resulted in an error:\n%s'%(first_error, error_msg))
else:
# 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)
#print('output=%s'%outputs)
#print('expected=%s'%expected)
#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)
else:
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)
else:
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:
self.strip_keys(output)
return cell
def transform_message(self, msg, expected):
"""
transform a message into something like the notebook
"""
SWAP_KEYS = {
"output_type": {
"pyout": "execute_result",
"pyerr": "error"
}
}
output = {
u"output_type": msg["msg_type"]
}
output.update(msg["content"])
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.