import sys

import unittest

import StringIO

import atexit

from test import test_support

class TestCase(unittest.TestCase):

    def setUp(self):

        s = StringIO.StringIO()

        self.save_stdout = sys.stdout

        self.save_stderr = sys.stderr

        sys.stdout = sys.stderr = self.subst_io = s

        self.save_handlers = atexit._exithandlers

        atexit._exithandlers = []

    def tearDown(self):

        sys.stdout = self.save_stdout

        sys.stderr = self.save_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):

        atexit.register(lambda: 1, 0, 0, (x for x in (1,2)), 0, 0)

        self.assertRaises(TypeError, atexit._run_exitfuncs)

    def test_order(self):

        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

        exfunc = sys.exitfunc

        sys.exitfunc = self.h1

        reload(atexit)

        try:

            atexit.register(self.h2)

            atexit._run_exitfuncs()

        finally:

            sys.exitfunc = exfunc

        self.assertEqual(self.subst_io.getvalue(), "h2\nh1\n")

    def test_raise(self):

        atexit.register(self.raise1)

        atexit.register(self.raise2)

        self.assertRaises(TypeError, atexit._run_exitfuncs)

    ### helpers

    def h1(self):

        print "h1"

    def h2(self):

        print "h2"

    def h3(self):

        print "h3"

    def h4(self, *args, **kwargs):

        print "h4", args, kwargs

    def raise1(self):

        raise TypeError

    def raise2(self):

        raise SystemError

def test_main():

    test_support.run_unittest(TestCase)

if __name__ == "__main__":

    test_main()