test_nbs.py

# We need a "test_" file with "test_" functions to make it easy to run with pytest


# couple of example "test_" functions

# import nbdev.test
# def test_run():
#     print('running nbdev.test.test_nb("20_models.ipynb") ...')
#     nbdev.test.test_nb('20_models.ipynb')

# import os
# def test_run():
#     print('running nbdev_test_nbs...')
#     os.system('nbdev_test_nbs')


# set-up a "before test" callback handler that will modify the notebook before it is run
import os
from pathlib import Path
import time
import glob
import nbformat
from nbdev.imports import Config,parallel
from nbdev.export import read_nb,find_default_export,is_export,split_flags_and_code
from nbdev.test import get_all_flags,NoExportPreprocessor

def before_test(nb):
    "callback that will import modules and run cells that are not exported"
    default_export=find_default_export(nb['cells'])
    exports = [is_export(c, default_export) for c in nb['cells']]
    imports = ''
    # exclude exported, notebook2script calls etc
    things_to_exclude = ['notebook2script']
    cells = [(i,c,e) for i,(c,e) in enumerate(zip(nb['cells'],exports)) if c['cell_type']=='code']
    for i,c,e in cells:
        if e: 
            c['cell_type']='exclude'             # if it's exported to the library, don't run as test
            for line in split_flags_and_code(c): # but we might still need to run import statements
                if 'import' in line: imports += f'{line}\n'
            continue 
        for thing_to_exclude in things_to_exclude: # TODO: is this too coarse? maybe just exclude specific lines?
            if thing_to_exclude in c['source']: 
                c['cell_type']='exclude'
                continue
    
    nb['cells'].insert(0,nbformat.v4.new_code_cell(imports))
    
    # import everything from modules written to by this notebook
    for export in {export[0] for export in exports if export}:
        export_parts=export.split('.')
        b=export_parts.pop()
        export_parts.insert(0,Config().lib_name)
        a='.'.join(export_parts)
        src=f"""
from {a} import {b}
for o in dir({b}):
    exec(f'from {a}.{b} import {{o}}')"""
        nb['cells'].insert(0,nbformat.v4.new_code_cell(src))
    return nb

# uncomment to see current nbdev behaviour
# i.e. use a before test callback that does nothing
# def before_test(nb): return nb

# If nbdev.test.test_nb knew to call our "before test" callback, the rest of this script could be just the following 3 lines
# def test_run():
#     from nbdev.cli import nbdev_test_nbs
#     nbdev_test_nbs.__wrapped__()

# until it does ... we need to duplicate a few chunks of nbdev
def _test_nb(fn, flags=None):
    "Execute tests in notebook in `fn` with `flags`"
    os.environ["IN_TEST"] = '1'
    if flags is None: flags = []
    try:
        nb = read_nb(fn)
        nb = before_test(nb) # <- THIS is the only change to nbdev code
        for f in get_all_flags(nb['cells']):
            if f not in flags: return
        ep = NoExportPreprocessor(flags, timeout=600, kernel_name='python3')
        pnb = nbformat.from_dict(nb)
        ep.preprocess(pnb)
    finally: os.environ.pop("IN_TEST")

def _test_one(fname, flags=None, verbose=True):
    print(f"testing: {fname}")
    start = time.time()
    try: 
        _test_nb(fname, flags=flags)
        return True,time.time()-start
    except Exception as e: 
        if "Kernel died before replying to kernel_info" in str(e):
            time.sleep(random.random())
            _test_one(fname, flags=flags)
        if verbose: print(f'Error in {fname}:\n{e}')
        return False,time.time()-start

def nbdev_test_nbs(fname=None,flags=None,n_workers=None,verbose=True,timing=False):
    """
    fname:Param("A notebook name or glob to convert", str)=None,
    flags:Param("Space separated list of flags", str)=None,
    n_workers:Param("Number of workers to use", int)=None,
    verbose:Param("Print errors along the way", bool)=True,
    timing:Param("Timing each notebook to see the ones are slow", bool)=False
    """
    "Test in parallel the notebooks matching `fname`, passing along `flags`"
    if flags is not None: flags = flags.split(' ')
    if fname is None: 
        files = [f for f in Path(Config(".", "planetarypy").config_path).glob('*.ipynb') if not f.name.startswith('_')]
    else: files = glob.glob(fname)
    files = [Path(f).absolute() for f in sorted(files)]
    if len(files)==1 and n_workers is None: n_workers=0
    # make sure we are inside the notebook folder of the project
    os.chdir(Config(".", "planetarypy").config_path)
    results = parallel(_test_one, files, flags=flags, verbose=verbose, n_workers=n_workers)
    passed,times = [r[0] for r in results],[r[1] for r in results]
    if all(passed): print("All tests are passing!")
    else:
        msg = "The following notebooks failed:\n"
        raise Exception(msg + '\n'.join([f.name for p,f in zip(passed,files) if not p]))
    if timing:
        for i,t in sorted(enumerate(times), key=lambda o:o[1], reverse=True): 
            print(f"Notebook {files[i].name} took {int(t)} seconds")

def test_run():
    # now we can "nbdev_test_nbs" and have our "before test" callback called
    # nbdev_test_nbs('00_core.ipynb') # Use this line to test a single notebook
    nbdev_test_nbs()