How to mock/monkeypatch during notebook execution

I have some notebooks that I want to execute programmatically in a pytest test. I can use papermill or ExecutePreprocessor from nbconvert and these both work great. But I also want to be able to patch various functions and classes that my notebook calls (i.e. source code that is not written in the notebook). Using the regular suite of patching/mocking tools doesn’t work I presume because the notebooks are not being run in pytest's environment.

Let’s consider a super simple example:

notebook.ipynb

from mymodule import bar
bar()

mymodule.py

def bar():
    print("bar")

test_module.py

import nbformat
from nbconvert.preprocessors import ExecutePreprocessor
import mymodule 

def test_foo(monkeypatch):

    monkeypatch.setattr(mymodule, "bar", lambda: print("foo"))

    with open("notebook.ipynb") as f:
        nb = nbformat.read(f, as_version=4)

    ep = ExecutePreprocessor(timeout=600, kernel_name='python3')
    ep.preprocess(nb)  # prints "bar"

The monkeypatching does not work here, we still print “bar” rather than “foo”. One way I have come up with is to modify the nb object by inserting the following code into the notebook (e.g. in the first cell) after loading but before executing:

import mymodule
import pytest

mp = pytest.MonkeyPatch()
mp.setattr(mymodule, "bar", lambda: print("foo"))

This works and we now print “foo” when the test runs but this is quite inelegant. My question is whether there is another nicer way to do this?

Have you seen testbook?
It says it supports injecting code and patching.

I learned of it in the Testing notebooks thread here. There may be other things to look into there.

1 Like

Thanks so much @fomightez! Testbook is just what I needed. Worth mentioning that my proposed solution is basically exactly what they are doing here, just injecting code into the notebook prior to execution - but it’s much easier to use a library that is providing you wrappers for these things than to do it yourself!

1 Like