pytest测试Jupyter Notebook返回值异常问题求助

Let's break down the problem and fix it step by step:

The Root Cause #

Your _notebook_run function uses subprocess.check_call(), which returns the exit code of the jupyter nbconvert command—not the actual output of your Notebook's code.

When jupyter nbconvert --execute runs successfully (even if your Notebook's code outputs something else), it will always return an exit code of 0. It only returns a non-zero exit code if the Notebook execution fails (e.g., your code throws an error). That's why your test keeps getting 0, no matter what your Notebook outputs.

How to Fix It #

To get the actual output from your Notebook, you need to capture and parse the execution results of the Notebook itself. Here's a reliable approach using nbformat (a library built for working with Jupyter Notebook files):

1. First, install the required dependency if you haven't already:

   pip install nbformat
   

2. Rewrite your helper function to execute the Notebook, save the executed version, and extract the output:

   import subprocess
   import nbformat
   import os
   
   def _get_notebook_output(path):
       # Define a path to save the executed Notebook (with outputs included)
       executed_notebook_path = "./executed_temp.ipynb"
       
       # Command to execute the Notebook and save the result
       args = [
           "jupyter", "nbconvert", "--execute", path,
           "--to", "notebook", "--output", executed_notebook_path,
           "--ExecutePreprocessor.timeout=60"  # Prevent hanging on long-running code
       ]
       
       # Run the command; this will throw an error if execution fails
       subprocess.check_call(args)
       
       # Read and parse the executed Notebook
       with open(executed_notebook_path, 'r', encoding='utf-8') as f:
           notebook = nbformat.read(f, as_version=4)
       
       # Clean up the temporary file
       os.remove(executed_notebook_path)
       
       # Extract the output from code cells (targets the last code cell's output here)
       for cell in reversed(notebook.cells):
           if cell.cell_type == "code" and cell.outputs:
               # Loop through outputs in the cell to find the text output
               for output in cell.outputs:
                   if "text" in output:
                       # Clean up the output (remove newlines) and convert to integer
                       return int(output["text"].strip())
       
       # Return None if no valid output was found
       return None
   

3. Update your test function to check the actual Notebook output:

   def test_notebook_output():
       notebook_output = _get_notebook_output('./file.ipynb')
       assert notebook_output == 1, f"Expected output 1, got {notebook_output}"
   

Key Notes #

• This method captures the actual text output from your Notebook's code cells. The example targets the last code cell's output (using reversed()), but you can adjust the logic to target a specific cell if needed.
• If your Notebook execution fails (e.g., code raises an exception), subprocess.check_call() will throw a subprocess.CalledProcessError, which will automatically fail your test—this helps catch unexpected errors in your Notebook.
• The temporary executed Notebook file is deleted right after parsing to avoid clutter.

内容的提问来源于stack exchange，提问作者president