Home | History | Annotate | Download | only in test 
      1 import sys
      2 import unittest
      3 import StringIO
      4 import atexit
      5 from imp import reload
      6 from test import test_support
      7 
      8 class 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 
     77 def test_main():
     78     test_support.run_unittest(TestCase)
     79 
     80 
     81 if __name__ == "__main__":
     82     test_main()
     83