172 lines
5.4 KiB
Python
172 lines
5.4 KiB
Python
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")
|
|
|
|
|