RE #27779 Add plotting model and tests

scripts/Engineering/gui/engineering_diffraction/tabs/fitting/plotting/plot_model.py

+# Mantid Repository : https://github.com/mantidproject/mantid
+#
+# Copyright © 2020 ISIS Rutherford Appleton Laboratory UKRI,
+#     NScD Oak Ridge National Laboratory, European Spallation Source
+#     & Institut Laue - Langevin
+# SPDX - License - Identifier: GPL - 3.0 +
+
+
+class FittingPlotModel(object):
+    def __init__(self):
+        self.plotted_workspaces = []
+
+    def get_plotted_workspaces(self):
+        return self.plotted_workspaces
+
+    def add_workspace_to_plot(self, ws, ax, plot_kwargs):
+        ax.plot(ws, **plot_kwargs)
+        self.plotted_workspaces.append(ws)
+
+    def remove_workspace_from_plot(self, ws, ax):
+        if ws in self.plotted_workspaces:
+            self._remove_workspace_from_plot(ws, ax)
+            self.plotted_workspaces.remove(ws)
+
+    @staticmethod
+    def _remove_workspace_from_plot(ws, ax):
+        ax.remove_workspace_artists(ws)
+
+    def remove_all_workspaces_from_plot(self, ax):
+        for ws in self.plotted_workspaces:
+            self._remove_workspace_from_plot(ws, ax)
+        # Should already be empty, but just to be sure.
+        self.plotted_workspaces.clear()

scripts/Engineering/gui/engineering_diffraction/tabs/fitting/plotting/test/test_plot_model.py

+# Mantid Repository : https://github.com/mantidproject/mantid
+#
+# Copyright © 2020 ISIS Rutherford Appleton Laboratory UKRI,
+#     NScD Oak Ridge National Laboratory, European Spallation Source
+#     & Institut Laue - Langevin
+# SPDX - License - Identifier: GPL - 3.0 +
+
+import unittest
+
+from mantid.py3compat import mock
+
+from Engineering.gui.engineering_diffraction.tabs.fitting.plotting import plot_model
+
+dir_path = "Engineering.gui.engineering_diffraction.tabs.fitting.plotting"
+
+
+class FittingPlotModelTest(unittest.TestCase):
+    def setUp(self):
+        self.model = plot_model.FittingPlotModel()
+
+    def test_adding_workspace_to_plot(self):
+        self.assertEqual([], self.model.plotted_workspaces)
+        ax = mock.MagicMock()
+
+        self.model.add_workspace_to_plot("mocked_ws", ax, {"linestyle": "x"})
+
+        self.assertEqual(["mocked_ws"], self.model.plotted_workspaces)
+        ax.plot.assert_called_once_with("mocked_ws", linestyle="x")
+
+    def test_removing_single_tracked_workspace_from_plot(self):
+        self.model.plotted_workspaces.append("mocked_ws")
+        ax = mock.MagicMock()
+
+        self.model.remove_workspace_from_plot("mocked_ws", ax)
+
+        self.assertEqual([], self.model.plotted_workspaces)
+        ax.remove_workspace_artists.assert_called_once_with("mocked_ws")
+
+    def test_removing_not_tracked_workspace_from_plot(self):
+        self.model.plotted_workspaces.append("mocked_ws")
+        ax = mock.MagicMock()
+
+        self.model.remove_workspace_from_plot("whatever", ax)
+
+        self.assertEqual(["mocked_ws"], self.model.plotted_workspaces)
+        ax.remove_workspace_artists.assert_not_called()
+
+    def test_removing_all_workspaces_from_plot(self):
+        self.model.plotted_workspaces.extend(["mocked_ws", "mock_ws_2"])
+        ax = mock.MagicMock()
+
+        self.model.remove_all_workspaces_from_plot(ax)
+
+        self.assertEqual([], self.model.plotted_workspaces)
+        ax.remove_workspace_artists.assert_any_call("mocked_ws")
+        ax.remove_workspace_artists.assert_any_call("mock_ws_2")
+        self.assertEqual(2, ax.remove_workspace_artists.call_count)
+
+
+if __name__ == '__main__':
+    unittest.main()