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()