Unit Testing Python data visualizations
By: Tristan Crockett
Tags: Data Visualization, Data Science, Matplotlib, Python, Unit Testing

Unit Testing Python code that interfaces with matplotlib can be challenging. How much is helpful to test, and what’s the best way of doing so?

This comes up a lot while doing development for Triage, a data science toolkit that, among other things, includes a decent amount of code that draws plots using matplotlib or other related Python libraries. These plots are often used in Jupyter notebooks, which have their own facilities to show the plots.

It doesn’t make a ton of sense to test visualization code in too much detail. Expectations about what a plot should look like are hard to express in a way that is easy to unit test. It usually makes sense to separate as much of the plot data generation from the plotting as possible, and test the plot data generation as you would any other generated data. But plot data is not a plot. We want to make sure that the plotting code works. How do we define ‘works’? One useful definition is if a test can answer the question:

Did we plot something?

This is a great start. Any code that is supposed to plot something should probably plot something, and hopefully it should be easy enough to figure this out automatically.

Approach one: patch

The approach I’ve used a lot in the past to do this is just to mock ‘plt.show’ and assert that it got called. An example:

This works, generally. Because ‘plt.show’ is generally called at the end of plotting code, if something breaks leading up to plt.show (often what you’re trying to test), this mock won’t get called and the test will fail.

There are some shortcomings, though:

What if plt.show isn’t called? For code that’s often called from a Jupyter notebook (or any other display engine that may be handling the showing on its own layer). You can try and make internal guidelines to enforce that all plotting code calls plt.show, but at that point you’re just testing for your internal guidelines, not that the code works.

What if plt.show is called, but not in the module you’re importing? Take seaborn plots, for example. They use matplotlib but wrap it, so you’re less likely to see plt.show in your code. What if somebody refactors a clunky matplotlib figure into a nice seaborn one? I don’t want my test to break in this case.

What if plt.show was called but without any content? This is not super likely, but still possible.

Approach two: plt.gcf

Luckily, matplotlib has another interface we can use to get around these shortcomings of patch: plt.gcf() (get current figure).

I’m not sure if this is how Jupyter’s inline plotting magic works, but it’s mighty useful. gcf returns a reference to the current (last plotted) figure. Better yet, it includes a number (plt.gcf().number) which lets you keep track of how many were plotted. We can inspect the current figure after our plotting code is run to make whatever level of assertions we want on it. So it doesn’t care if plt.show was called, and works as long as somewhere the figures got created, even some higher-level library that internally uses matplotlib.

When I first found gcf, it was tempting to “reset” the plot number before calling the code to try and assert what the number of the current figure should be exactly. I tried to “clear” the plot before doing so, with plt.close():

This worked….locally. However, when I pushed to Travis the ‘close’ didn’t seem to do anything and the value of gcf().number was affected by other test cases. I assume this has to do with my laptop having an actual display/window and Travis being headless. Maybe there’s a better way to do the ‘clearing’, but I wasn’t about to get into the weeds with displays, so instead I came up with an alternate way to make it deterministic:

assert that the number of figures is higher than when you started the test

This works the same on Travis as on my laptop. Another difference is that I’m not asserting a specific number of figures plotted. Whether or not you want to do this in your code (ultimately I think it may get too deeply into how matplotlib is handling things internally), you could using this method. But this by itself is a handy plot assertion tool. I even ended up making it into a context manager-aware method for easy reuse:

What are the possible shortcomings of this method?

We’re asserting that whatever code we are calling to plot creates matplotlib figures. Is this brittle? Theoretically. It definitely sounds like it violates the “don’t test the implementation” guideline. But in all reality, I think that

matplotlib is a common backend that many nice frontends that i may switch between will use

is a safer assumption than

seaborn is the best frontend and i will always use it and maybe it will switch backends away from matplotlib

My team redoing a chart into a new, non-matplotlib-backed plotting library is certainly possible, but not very likely. So I’m fairly comfortable with this as a reasonably-generic Python data visualization assertion. What do you think? Are there other shortcomings I’m missing, or better ways to test Python plotting?