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

I am not sure how matplotlib interacts with pytest however it seems like you need to use fixtures in order to achieve this. You also need to create some sort of assert statement in your test that will signal for the test to gracefully tear down your fixture. Some thing like this should achieve the desired results.

import pytest 
import matplotlib.pyplot as plt 

@pytest.fixture(scope='function')
def plot_fn():
    def _plot(points):
        plt.plot(points)
        yield plt.show()
        plt.close('all')
    return _plot


def test_plot_fn(plot_fn):
    points = [1, 2, 3]
    plot_fn(points)
    assert True

If you want to simply monkeypatch the show behavior I would do it like shown below.

def test_plot_fn(monkeypatch):
    monkeypatch.setattr(plt, 'show', lambda: None)
    plot_fn()

Answer 2

I can think about two options:

1. use the plt.show(block=False) and still close the figure when test ends.

2. in test instead of plt.show save the plot to a file plt.savefig. This way you can also compare the created file to a test ref and also test that you plot is the same. Of course for this option you'll have to let your function "know" it is in test mode by optional parameter, env variable etc.

Answer 3

This one-line decorator will patch the plt.show() behavior and suppress the wait/hang part of showing the pyplot figure during a unit test, like this:

import pytest
import unittest
try:
    # python 3.4+ should use builtin unittest.mock not mock package
    from unittest.mock import patch
except ImportError:
    from mock import patch


@patch("methylcheck.qc_plot.plt.show")
def test_plot_fn(mock_this):
   plot_fn()
   plt.close('all')

Note: you have to patch the show() function within the namespace path of the module/package you are testing, so in this case, the methylcheck package has a qc_plot file that imports pyplot in the usual way (import matplotlib.pyplot as plt) and THAT import gets patched.