1import sys 2import unittest 3import StringIO 4import atexit 5from imp import reload 6from test import test_support 7 8class TestCase(unittest.TestCase): 9 def setUp(self): 10 s = StringIO.StringIO() 11 self.save_stdout = sys.stdout 12 self.save_stderr = sys.stderr 13 sys.stdout = sys.stderr = self.subst_io = s 14 self.save_handlers = atexit._exithandlers 15 atexit._exithandlers = [] 16 17 def tearDown(self): 18 sys.stdout = self.save_stdout 19 sys.stderr = self.save_stderr 20 atexit._exithandlers = self.save_handlers 21 22 def test_args(self): 23 atexit.register(self.h1) 24 atexit.register(self.h4) 25 atexit.register(self.h4, 4, kw="abc") 26 atexit._run_exitfuncs() 27 self.assertEqual(self.subst_io.getvalue(), 28 "h4 (4,) {'kw': 'abc'}\nh4 () {}\nh1\n") 29 30 def test_badargs(self): 31 atexit.register(lambda: 1, 0, 0, (x for x in (1,2)), 0, 0) 32 self.assertRaises(TypeError, atexit._run_exitfuncs) 33 34 def test_order(self): 35 atexit.register(self.h1) 36 atexit.register(self.h2) 37 atexit.register(self.h3) 38 atexit._run_exitfuncs() 39 self.assertEqual(self.subst_io.getvalue(), "h3\nh2\nh1\n") 40 41 def test_sys_override(self): 42 # be sure a preset sys.exitfunc is handled properly 43 exfunc = sys.exitfunc 44 sys.exitfunc = self.h1 45 reload(atexit) 46 try: 47 atexit.register(self.h2) 48 atexit._run_exitfuncs() 49 finally: 50 sys.exitfunc = exfunc 51 self.assertEqual(self.subst_io.getvalue(), "h2\nh1\n") 52 53 def test_raise(self): 54 atexit.register(self.raise1) 55 atexit.register(self.raise2) 56 self.assertRaises(TypeError, atexit._run_exitfuncs) 57 58 ### helpers 59 def h1(self): 60 print "h1" 61 62 def h2(self): 63 print "h2" 64 65 def h3(self): 66 print "h3" 67 68 def h4(self, *args, **kwargs): 69 print "h4", args, kwargs 70 71 def raise1(self): 72 raise TypeError 73 74 def raise2(self): 75 raise SystemError 76 77def test_main(): 78 test_support.run_unittest(TestCase) 79 80 81if __name__ == "__main__": 82 test_main() 83