How can I write unit tests against code that uses Matplotlib?

To write unit test cases against a code, we can consider a plot that takes an array as x points and plot it as y=x^2. While testing, we would extract y_data for x data points.−

Steps

• Create a method, i.e., plot_sqr_curve(x) to plot x and x^2 using plot() method and return the plot.
• To test, use unittest.TestCase.
• Write test_curve_sqr_plot() method that includes the following statements.
  • Create data points for x to plot the curve.
  • Using the above x data points, create y data points.
  • Using x and y data points, plot the curve.
  • Using pt (from step 5), extract x and y data.
  • Check whether the given expression is true or not.

Example

import unittest
import numpy as np
from matplotlib import pyplot as plt
def plot_sqr_curve(x):
   """
   Plotting x points with y = x^2.
   """
return plt.plot(x, np.square(x))
class TestSqrCurve(unittest.TestCase):
def test_curve_sqr_plot(self):
   x = np.array([1, 3, 4])
   y = np.square(x)
   pt, = plot_sqr_curve(x)
   y_data = pt.get_data()[1]
   x_data = pt.get_data()[0]
   self.assertTrue((y == y_data).all())
   self.assertTrue((x == x_data).all())
if __name__ == '__main__':
   unittest.main()

Output

Ran 1 test in 1.587s
OK