Add tests for figure update on workspace replacement

Refs #24000

Framework/PythonInterface/test/python/mantid/plots/plots__init__Test.py

@@ -9,10 +9,11 @@ from __future__ import (absolute_import, division, print_function)
 import matplotlib
 matplotlib.use('AGG')
 import matplotlib.pyplot as plt
+from matplotlib.container import ErrorbarContainer
 import numpy as np
 import unittest
 
-from mantid.simpleapi import CreateWorkspace, DeleteWorkspace
+from mantid.simpleapi import CreateWorkspace, CreateSampleWorkspace, DeleteWorkspace
 
 
 class Plots__init__Test(unittest.TestCase):
@@ -35,11 +36,89 @@ class Plots__init__Test(unittest.TestCase):
     def tearDownClass(cls):
         DeleteWorkspace('ws2d_histo')
 
-    def test_1d_plots(self):
-        fig, ax = plt.subplots(subplot_kw={'projection': 'mantid'})
-        # ax.plot(self.ws2d_histo, 'rs', specNum=1)
-        ax.plot(self.ws2d_histo, specNum=2, linewidth=6)
-        ax.plot(np.arange(10), np.arange(10), 'bo-')
+    def setUp(self):
+        self.fig, self.ax = plt.subplots(subplot_kw={'projection': 'mantid'})
+
+    def tearDown(self):
+        plt.close('all')
+        self.fig, self.ax = None, None
+
+    def test_line2d_plots(self):
+        self.ax.plot(self.ws2d_histo, specNum=2, linewidth=6)
+        self.ax.plot(np.arange(10), np.arange(10), 'bo-')
+
+    def test_errorbar_plots(self):
+        self.ax.errorbar(self.ws2d_histo, specNum=2, linewidth=6)
+        self.ax.errorbar(np.arange(10), np.arange(10), 0.1*np.ones((10,)), fmt='bo-')
+
+    def test_imshow(self):
+        self.ax.imshow(self.ws2d_histo)
+
+    def test_pcolor(self):
+        self.ax.pcolor(self.ws2d_histo)
+
+    def test_pcolorfast(self):
+        self.ax.pcolorfast(self.ws2d_histo)
+
+    def test_pcolormesh(self):
+        self.ax.pcolormesh(self.ws2d_histo)
+
+    def test_remove_workspace_artist_for_known_workspace_removes_plot(self):
+        self.ax.plot(self.ws2d_histo, specNum=2, linewidth=6)
+        is_empty = self.ax.remove_workspace_artists(self.ws2d_histo)
+        self.assertEqual(True, is_empty)
+        self.assertEqual(0, len(self.ax.lines))
+
+    def test_remove_workspace_artist_for_unknown_workspace_does_nothing(self):
+        self.ax.plot(self.ws2d_histo, specNum=2, linewidth=6)
+        unknown_ws = CreateSampleWorkspace()
+        self.ax.remove_workspace_artists(unknown_ws)
+        self.assertEqual(1, len(self.ax.lines))
+
+    def test_remove_workspace_artist_for_removes_only_specified_workspace(self):
+        second_ws = CreateSampleWorkspace()
+        line_ws2d_histo = self.ax.plot(self.ws2d_histo, specNum=2, linewidth=6)[0]
+        line_second_ws = self.ax.plot(second_ws, specNum=5)[0]
+        self.assertEqual(2, len(self.ax.lines))
+
+        self.ax.remove_workspace_artists(self.ws2d_histo)
+        self.assertEqual(1, len(self.ax.lines))
+        self.assertTrue(line_ws2d_histo not in self.ax.lines)
+        self.assertTrue(line_second_ws in self.ax.lines)
+        DeleteWorkspace(second_ws)
+
+    def test_replace_workspace_data_plot(self):
+        plot_data = CreateWorkspace(DataX=[10, 20, 30, 10, 20, 30, 10, 20, 30],
+                                    DataY=[3, 4, 5, 3, 4, 5],
+                                    DataE=[1, 2, 3, 4, 1, 1],
+                                    NSpec=3)
+        line_ws2d_histo = self.ax.plot(plot_data, specNum=2, color='r')[0]
+        plot_data = CreateWorkspace(DataX=[20, 30, 40, 20, 30, 40, 20, 30, 40],
+                                    DataY=[3, 4, 5, 3, 4, 5],
+                                    DataE=[1, 2, 3, 4, 1, 1],
+                                    NSpec=3)
+        self.ax.replace_workspace_artists(plot_data)
+        self.assertAlmostEqual(25, line_ws2d_histo.get_xdata()[0])
+        self.assertAlmostEqual(35, line_ws2d_histo.get_xdata()[-1])
+        self.assertEquals('r', line_ws2d_histo.get_color())
+
+    def test_replace_workspace_data_errorbar(self):
+        eb_data = CreateWorkspace(DataX=[10, 20, 30, 10, 20, 30, 10, 20, 30],
+                                  DataY=[3, 4, 5, 3, 4, 5],
+                                  DataE=[1, 2, 3, 4, 1, 1],
+                                  NSpec=3)
+        self.ax.errorbar(eb_data, specNum=2, color='r')
+        eb_data = CreateWorkspace(DataX=[20, 30, 40, 20, 30, 40, 20, 30, 40],
+                                     DataY=[3, 4, 5, 3, 4, 5],
+                                     DataE=[.1, .2, .3, .4, .1, .1],
+                                     NSpec=3)
+        self.ax.replace_workspace_artists(eb_data)
+        self.assertEqual(1, len(self.ax.containers))
+        eb_container = self.ax.containers[0]
+        self.assertTrue(isinstance(eb_container, ErrorbarContainer))
+        self.assertAlmostEqual(25, eb_container[0].get_xdata()[0])
+        self.assertAlmostEqual(35, eb_container[0].get_xdata()[-1])
+        self.assertEquals('r', eb_container[0].get_color())
 
     def test_3d_plots(self):
         fig = plt.figure()