Pytest test function that creates plots

Question

I have several functions that create plots, which I use in Jupyter notebooks to visualise data.

I want to create basic tests for these, checking that they still run without erroring on various inputs if I make changes. However, if I call these functions using pytest, creating the plots causes the program to hang until I manually minimise the plot.

import pytest 
import matplotlib.pyplot as plt 

def plot_fn():
    plt.plot([1,2,3])
    plt.show()


def test_plot_fn():
   plot_fn()

How can I test that functions like 'plot_fn' run without erroring using Pytest? I tried the following, but it doesn't work, I think because plt.show() causes the script to hang, and so not reach plt.close('all').

def test_plot_fn():
   plot_fn()
   plt.close('all')

I'm happy to change the behaviour of my plotting function, for example to return the plt object?

Answer 1

This works.

from unittest.mock import patch 
import pytest 
import matplotlib.pyplot as plt 

def plot_fn():
    plt.plot([1,2,3])
    plt.show()

@patch("matplotlib.pyplot.show")
def test_plot_fn(mock_show):
    plot_fn()

Based on this answer (possible duplicate) Turn off graphs while running unittests