Turn velocity plot test into a figure test

punchbowl/level3/tests/test_velocity.py

@@ -74,15 +74,14 @@ def test_shape_matching(synthetic_data):
     assert result.data.shape[0] == len(ycens)
 
 
+@pytest.mark.mpl_image_compare(style="default")
 def test_wind_plot(synthetic_data):
     """Tests that wind plots are generated and  output to file."""
     files = synthetic_data
     ycens = np.arange(7, 14.5, 0.5)
     result = track_velocity(files, ycens=ycens)
 
-    plot_flow_map(THIS_DIRECTORY / 'wind_map.png', result)
-
-    assert os.path.exists(THIS_DIRECTORY / 'wind_map.png')
+    return plot_flow_map(None, result)
 
 
 def test_no_nans_or_negatives(synthetic_data):

punchbowl/level3/velocity.py

@@ -365,21 +365,26 @@
     return avg_speeds, sigmas
 
 
-def plot_flow_map(filename: str, data: NDCube, cmap: str = "magma") -> None:
+def plot_flow_map(filename: str | None, data: NDCube, cmap: str = "magma") -> None:
     """
     Plot polar maps of the radial flows.
 
     Parameters
     ----------
     filename: str
-        Output plot filename
+        Output plot filename. If None, the figure is not saved out.
 
     data: NDCube
         Flow tracking data NDCube
 
     cmap : str, optional
         Colormap for the plot (default is 'magma')
 
+    Returns
+    -------
+    fig
+        The generated Matplotlib Figure
+
     """
     speeds = data.data
     sigmas = data.uncertainty.array
@@ -418,7 +423,9 @@
 
     cbar_ax = fig.add_axes([0.11, 0.2, 0.8, 0.03])
     plt.colorbar(mapper, cax=cbar_ax, orientation="horizontal").ax.set_xlabel("Speed (km/s)")
-    plt.savefig(filename)
+    if filename is not None:
+        plt.savefig(filename)
+    return fig
 
 
 def track_velocity(files: list[str],

pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+addopts = --mpl