"""
Execute each notebook as a test, reporting an error if any cell throws an exception.
Adapted from https://gist.github.com/minrk/2620876.
"""

import os

import nbformat
import pytest
from jupyter_client import KernelManager

test_dir = os.path.dirname(os.path.abspath(__file__))
examples = [fn for fn in os.listdir(test_dir) if fn.endswith('.ipynb')]


@pytest.fixture(params=examples)
def example_fn(request):
    if 'openmm' in request.param:
        try:
            from simtk.openmm import app
        except ImportError:
            pytest.skip("Openmm required for example notebook `{}`".format(request.param))

    cwd = os.path.abspath('.')
    os.chdir(test_dir)
    yield request.param
    os.chdir(cwd)


def test_example_notebook(example_fn):
    with open(example_fn) as f:
        nb = nbformat.reads(f.read(), nbformat.NO_CONVERT)
    run_notebook(nb)


def run_notebook(nb):
    km = KernelManager()
    km.start_kernel(stderr=open(os.devnull, 'w'))
    kc = km.client()
    kc.start_channels()
    shell = kc.shell_channel
    # simple ping:
    kc.execute("pass")
    shell.get_msg()

    failures = 0
    for cell in nb.cells:
        if cell.cell_type != 'code':
            continue
        kc.execute(cell.source)
        # wait for finish, maximum 20s
        reply = shell.get_msg(timeout=60)['content']
        if reply['status'] == 'error':
            failures += 1
            print("\nFAILURE:")
            print('\n'.join(reply['traceback']))
            print()

    kc.stop_channels()
    km.shutdown_kernel()
    del km
    if failures > 0:
        raise Exception()

@@ -142,17 +142,65 @@
 %22%22%22%0A
-%0Aimport o
+from __future__ import print_function%0Aimport os%0Aimport sy
 s%0A%0Ai
@@ -268,16 +268,99 @@
 lManager
+%0Afrom queue import Empty%0A%0AFLAKEY_LIST = %5B'centroids.ipynb'%5D%0ATIMEOUT = 60  # seconds
 %0A%0Atest_d
@@ -419,18 +419,78 @@
 ples = %5B
-fn
+pytest.mark.flaky(fn) if fn in FLAKEY_LIST else fn%0A           
  for fn 
@@ -790,16 +790,33 @@
 .format(
+%0A                
 request.
@@ -1442,16 +1442,33 @@
 source)%0A
+        try:%0A    
         
@@ -1490,20 +1490,23 @@
 sh, 
-maximum 20s%0A
+w/ timeout%0A    
     
@@ -1539,18 +1539,23 @@
 timeout=
-60
+TIMEOUT
 )%5B'conte
@@ -1559,16 +1559,209 @@
 ntent'%5D%0A
+        except Empty:%0A            raise Exception(%0A                'Timeout (%25.1f) when executing the following %25s cell: %22%25s%22' %25%0A                (TIMEOUT, cell.cell_type, cell.source.strip()))%0A
         
@@ -1847,16 +1847,33 @@
 AILURE:%22
+, file=sys.stderr
 )%0A      
@@ -1913,16 +1913,33 @@
 eback'%5D)
+, file=sys.stderr
 )%0A      
@@ -1950,16 +1950,31 @@
   print(
+file=sys.stderr
 )%0A%0A    k

Update test_examples.py