Home > Software design >  Testing a Jupyter Notebook
Testing a Jupyter Notebook

Time:01-19

I am trying to come up with a method to test a number of Jupyter notebooks. A test should run when a new notebook is implemented in a Github branch and submitted for a pull request. The tests are not that complicated, they are mostly just testing if the notebook runs end-to-end and without any errors, and maybe a few asserts. However:

  • There are certain calls in some cells that need to be mocked, e.g. a call to download the data from a database.
  • There may be some magic cells in the notebooks which run a pip command or something else.

I am open to use any testing library, such as 'pytest' or unittest, although pytest is preferred.

I looked at a few libraries for testing notebooks such as enter image description here

The trick is to inject a cell before my call to bigquery.Client and mock it:

from testbook import testbook

@testbook('./my_notebook.ipynb')
def test_get_details(tb):
    tb.inject(
        """
        import mock
        mock_client = mock.MagicMock()
        mock_df = pd.DataFrame()
        mock_df['week'] = range(10)
        mock_df['count'] = 5
        p1 = mock.patch.object(bigquery, 'Client', return_value=mock_client)
        mock_client.query().result().to_dataframe.return_value = mock_df
        p1.start()
        """,
        before=2,
        run=False
    )
    tb.execute()
    dataframe = tb.get('dataframe')
    assert dataframe.shape == (10, 2)

    x = tb.get('x')
    assert x == 7
  •  Tags:  
  • Related