测试 Jupyter Notebook

Question

我正在尝试想出一种方法来测试多个 Jupyter 笔记本。当在 Github 分支中实现新笔记本并提交拉取请求时，应该运行测试。测试并不那么复杂，它们主要只是测试笔记本是否端到端运行并且没有任何错误，也许还有一些断言。然而：

• 某些单元格中存在某些需要模拟的调用，例如从数据库下载数据的调用。
• 笔记本中可能有一些神奇的单元可以运行 pip 命令或其他命令。

我愿意使用任何测试库，例如“pytest”或unittest，尽管pytest这是首选。

我查看了一些用于测试笔记本的库，例如nbmake、treon和testbook，但我无法使它们工作。我还尝试将笔记本转换为 python 文件，但神奇单元被转换为get_ipython().run_cell_magic(...)调用，这成为一个问题，因为pytest使用 python 而不是 ipython，并且get_ipython()仅在 ipython 中可用。

所以，我想知道考虑到所有这些，测试 Jupyter Notebook 的好方法是什么。任何帮助表示赞赏。

Answer 1

这是我自己使用testbook 的解决方案。假设我有一个名为my_notebook.ipynb以下内​​容的笔记本：

诀窍是在调用之前注入一个单元格bigquery.Client并模拟它：

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