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