TST: Added tests for remove of legends on figure and axes objects

cimarronm · cimarronm · commit e7f051f31e08 · 2014-02-09T13:22:30.000-08:00
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -5,7 +5,7 @@
 from six.moves import xrange
 
 import mock
-from nose.tools import assert_equal
+from nose.tools import assert_equal, assert_is
 import numpy as np
 
 from matplotlib.testing.decorators import image_comparison, cleanup
@@ -113,6 +113,19 @@ def test_legend_expand():
         ax.legend(loc=3, mode=mode, ncol=2)
 
 
+@cleanup
+def test_legend_remove():
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 1, 1)
+    lines = ax.plot(range(10))
+    leg = fig.legend(lines, "test")
+    leg.remove()
+    assert_equal(fig.legends, [])
+    leg = ax.legend("test")
+    leg.remove()
+    assert_is(ax.get_legend(), None)
+
+
 class TestLegendFunction(object):
     # Tests the legend function on the Axes and pyplot.
 
