add figure to plot test case

pystencils_tests/test_stencils.py

@@ -2,6 +2,7 @@ import pystencils as ps
 import sympy as sp
 
 from pystencils.stencil import coefficient_list, plot_expression
+import pystencils.plot as plt
 
 
 def test_coefficient_list():
@@ -9,7 +10,8 @@ def test_coefficient_list():
     expr = 2 * f[1] + 3 * f[-1]
     coff = coefficient_list(expr)
     assert coff == [3, 0, 2]
-    plot_expression(expr, matrix_form=True)
+    figure = plt.figure()
+    plot_expression(expr, matrix_form=True, figure=figure)
 
     f = ps.fields("f: double[3D]")
     expr = 2 * f[1, 0, 0] + 3 * f[0, -1, 0]
@@ -22,9 +24,11 @@ def test_coefficient_list():
 
     # in 3D plot only works if there are entries on every of the three 2D planes. In the above examples z-1 was empty
     expr = 2 * f[1, 0, 0] + 1 * f[0, -1, 0] + 1 * f[0, 0, 1] + f[0, 0, -1]
-    plot_expression(expr)
+    figure = plt.figure()
+    plot_expression(expr, figure=figure)
 
 
 def test_plot_expression():
     f = ps.fields("f: double[2D]")
-    plot_expression(2 * f[1, 0] + 3 * f[0, -1], matrix_form=True)
+    figure = plt.figure()
+    plot_expression(2 * f[1, 0] + 3 * f[0, -1], matrix_form=True, figure=figure)