backport

Lib/test/test_atexit.py

@@ -5,86 +5,52 @@ import atexit
 from test import test_support
 
 class TestCase(unittest.TestCase):
-    def test_args(self):
-        # be sure args are handled properly
+    def setUp(self):
         s = StringIO.StringIO()
-        sys.stdout = sys.stderr = s
-        save_handlers = atexit._exithandlers
+        sys.stdout = sys.stderr = self.subst_io = s
+        self.save_handlers = atexit._exithandlers
         atexit._exithandlers = []
-        try:
-            atexit.register(self.h1)
-            atexit.register(self.h4)
-            atexit.register(self.h4, 4, kw="abc")
-            atexit._run_exitfuncs()
-        finally:
-            sys.stdout = sys.__stdout__
-            sys.stderr = sys.__stderr__
-            atexit._exithandlers = save_handlers
-        self.assertEqual(s.getvalue(), "h4 (4,) {'kw': 'abc'}\nh4 () {}\nh1\n")
+
+    def tearDown(self):
+        sys.stdout = sys.__stdout__
+        sys.stderr = sys.__stderr__
+        atexit._exithandlers = self.save_handlers
+
+    def test_args(self):
+        atexit.register(self.h1)
+        atexit.register(self.h4)
+        atexit.register(self.h4, 4, kw="abc")
+        atexit._run_exitfuncs()
+        self.assertEqual(self.subst_io.getvalue(),
+                         "h4 (4,) {'kw': 'abc'}\nh4 () {}\nh1\n")
 
     def test_badargs(self):
-        s = StringIO.StringIO()
-        sys.stdout = sys.stderr = s
-        save_handlers = atexit._exithandlers
-        atexit._exithandlers = []
-        try:
-            atexit.register(lambda: 1, 0, 0, (x for x in (1,2)), 0, 0)
-            self.assertRaises(TypeError, atexit._run_exitfuncs)
-        finally:
-            sys.stdout = sys.__stdout__
-            sys.stderr = sys.__stderr__
-            atexit._exithandlers = save_handlers
+        atexit.register(lambda: 1, 0, 0, (x for x in (1,2)), 0, 0)
+        self.assertRaises(TypeError, atexit._run_exitfuncs)
 
     def test_order(self):
-        # be sure handlers are executed in reverse order
-        s = StringIO.StringIO()
-        sys.stdout = sys.stderr = s
-        save_handlers = atexit._exithandlers
-        atexit._exithandlers = []
-        try:
-            atexit.register(self.h1)
-            atexit.register(self.h2)
-            atexit.register(self.h3)
-            atexit._run_exitfuncs()
-        finally:
-            sys.stdout = sys.__stdout__
-            sys.stderr = sys.__stderr__
-            atexit._exithandlers = save_handlers
-        self.assertEqual(s.getvalue(), "h3\nh2\nh1\n")
+        atexit.register(self.h1)
+        atexit.register(self.h2)
+        atexit.register(self.h3)
+        atexit._run_exitfuncs()
+        self.assertEqual(self.subst_io.getvalue(), "h3\nh2\nh1\n")
 
     def test_sys_override(self):
         # be sure a preset sys.exitfunc is handled properly
-        save_handlers = atexit._exithandlers
-        atexit._exithandlers = []
         exfunc = sys.exitfunc
         sys.exitfunc = self.h1
         reload(atexit)
-        s = StringIO.StringIO()
-        sys.stdout = sys.stderr = s
         try:
             atexit.register(self.h2)
             atexit._run_exitfuncs()
         finally:
-            sys.stdout = sys.__stdout__
-            sys.stderr = sys.__stderr__
-            atexit._exithandlers = save_handlers
             sys.exitfunc = exfunc
-        self.assertEqual(s.getvalue(), "h2\nh1\n")
+        self.assertEqual(self.subst_io.getvalue(), "h2\nh1\n")
 
     def test_raise(self):
-        # be sure raises are handled properly
-        s = StringIO.StringIO()
-        sys.stdout = sys.stderr = s
-        save_handlers = atexit._exithandlers
-        atexit._exithandlers = []
-        try:
-            atexit.register(self.raise1)
-            atexit.register(self.raise2)
-            self.assertRaises(TypeError, atexit._run_exitfuncs)
-        finally:
-            sys.stdout = sys.__stdout__
-            sys.stderr = sys.__stderr__
-            atexit._exithandlers = save_handlers
+        atexit.register(self.raise1)
+        atexit.register(self.raise2)
+        self.assertRaises(TypeError, atexit._run_exitfuncs)
 
     ### helpers
     def h1(self):