add CleanupTestCase, a unittest.TestCase with cleanup

toddrjen · toddrjen · commit 6cf4d84c52c7 · 2013-10-22T22:38:32.000+02:00
diff --git a/lib/matplotlib/testing/decorators.py b/lib/matplotlib/testing/decorators.py
@@ -18,6 +18,7 @@
 from matplotlib.testing.compare import comparable_formats, compare_images, \
      make_test_filename
 import warnings
+import unittest
 
 def knownfailureif(fail_condition, msg=None, known_exception_class=None ):
     """
@@ -58,6 +59,7 @@ def failer(*args, **kwargs):
         return nose.tools.make_decorator(f)(failer)
     return known_fail_decorator
 
+
 class CleanupTest(object):
     @classmethod
     def setup_class(cls):
@@ -71,11 +73,30 @@ def teardown_class(cls):
 
         matplotlib.units.registry.clear()
         matplotlib.units.registry.update(cls.original_units_registry)
-        warnings.resetwarnings() #reset any warning filters set in tests
+        warnings.resetwarnings()  # reset any warning filters set in tests
 
     def test(self):
         self._func()
 
+
+class CleanupTestCase(unittest.TestCase):
+    '''A wrapper for unittest.TestCase that includes cleanup operations'''
+    @classmethod
+    def setUpClass(cls):
+        import matplotlib.units
+        cls.original_units_registry = matplotlib.units.registry.copy()
+
+    @classmethod
+    def tearDownClass(cls):
+        plt.close('all')
+
+        matplotlib.tests.setup()
+
+        matplotlib.units.registry.clear()
+        matplotlib.units.registry.update(cls.original_units_registry)
+        warnings.resetwarnings()  # reset any warning filters set in tests
+
+
 def cleanup(func):
     name = func.__name__
     func = staticmethod(func)
