From f68a4afce3f9c81ed146324f6215bf57ada17167 Mon Sep 17 00:00:00 2001
From: Pedro Bressan <phmbressan@usp.br>
Date: Fri, 8 Nov 2024 15:48:19 +0100
Subject: [PATCH] TST: implement testing of show or save plots.

---
 tests/unit/test_plots.py | 53 +++++++++++++++++++++++++++++++++++++++-
 1 file changed, 52 insertions(+), 1 deletion(-)

diff --git a/tests/unit/test_plots.py b/tests/unit/test_plots.py
index cd35f8d11..6d7c0cce2 100644
--- a/tests/unit/test_plots.py
+++ b/tests/unit/test_plots.py
@@ -1,8 +1,11 @@
-from unittest.mock import patch
+import os
+from unittest.mock import MagicMock, patch
 
 import matplotlib.pyplot as plt
+import pytest
 
 from rocketpy.plots.compare import Compare
+from rocketpy.plots.plot_helpers import show_or_save_fig, show_or_save_plot
 
 
 @patch("matplotlib.pyplot.show")
@@ -38,3 +41,51 @@ def test_compare(mock_show, flight_calisto):  # pylint: disable=unused-argument
     )
 
     assert isinstance(fig, plt.Figure)
+
+
+@patch("matplotlib.pyplot.show")
+@pytest.mark.parametrize("filename", [None, "test.png"])
+def test_show_or_save_plot(mock_show, filename):
+    """This test is to check if the show_or_save_plot function is
+    working properly.
+
+    Parameters
+    ----------
+    mock_show :
+        Mocks the matplotlib.pyplot.show() function to avoid showing
+        the plots.
+    filename : str
+        Name of the file to save the plot. If None, the plot will be
+        shown instead.
+    """
+    plt.subplots()
+    show_or_save_plot(filename)
+
+    if filename is None:
+        mock_show.assert_called_once()
+    else:
+        assert os.path.exists(filename)
+        os.remove(filename)
+
+
+@pytest.mark.parametrize("filename", [None, "test.png"])
+def test_show_or_save_fig(filename):
+    """This test is to check if the show_or_save_fig function is
+    working properly.
+
+    Parameters
+    ----------
+    filename : str
+        Name of the file to save the plot. If None, the plot will be
+        shown instead.
+    """
+    fig, _ = plt.subplots()
+
+    fig.show = MagicMock()
+    show_or_save_fig(fig, filename)
+
+    if filename is None:
+        fig.show.assert_called_once()
+    else:
+        assert os.path.exists(filename)
+        os.remove(filename)