1import sys 2import unittest 3import io 4import atexit 5from test import support 6 7### helpers 8def h1(): 9 print("h1") 10 11def h2(): 12 print("h2") 13 14def h3(): 15 print("h3") 16 17def h4(*args, **kwargs): 18 print("h4", args, kwargs) 19 20def raise1(): 21 raise TypeError 22 23def raise2(): 24 raise SystemError 25 26 27class GeneralTest(unittest.TestCase): 28 29 def setUp(self): 30 self.save_stdout = sys.stdout 31 self.save_stderr = sys.stderr 32 self.stream = io.StringIO() 33 sys.stdout = sys.stderr = self.stream 34 atexit._clear() 35 36 def tearDown(self): 37 sys.stdout = self.save_stdout 38 sys.stderr = self.save_stderr 39 atexit._clear() 40 41 def test_args(self): 42 # be sure args are handled properly 43 atexit.register(h1) 44 atexit.register(h4) 45 atexit.register(h4, 4, kw="abc") 46 atexit._run_exitfuncs() 47 48 self.assertEqual(self.stream.getvalue(), 49 "h4 (4,) {'kw': 'abc'}\nh4 () {}\nh1\n") 50 51 def test_badargs(self): 52 atexit.register(lambda: 1, 0, 0, (x for x in (1,2)), 0, 0) 53 self.assertRaises(TypeError, atexit._run_exitfuncs) 54 55 def test_order(self): 56 # be sure handlers are executed in reverse order 57 atexit.register(h1) 58 atexit.register(h2) 59 atexit.register(h3) 60 atexit._run_exitfuncs() 61 62 self.assertEqual(self.stream.getvalue(), "h3\nh2\nh1\n") 63 64 def test_raise(self): 65 # be sure raises are handled properly 66 atexit.register(raise1) 67 atexit.register(raise2) 68 69 self.assertRaises(TypeError, atexit._run_exitfuncs) 70 71 def test_raise_unnormalized(self): 72 # Issue #10756: Make sure that an unnormalized exception is 73 # handled properly 74 atexit.register(lambda: 1 / 0) 75 76 self.assertRaises(ZeroDivisionError, atexit._run_exitfuncs) 77 self.assertIn("ZeroDivisionError", self.stream.getvalue()) 78 79 def test_print_tracebacks(self): 80 # Issue #18776: the tracebacks should be printed when errors occur. 81 def f(): 82 1/0 # one 83 def g(): 84 1/0 # two 85 def h(): 86 1/0 # three 87 atexit.register(f) 88 atexit.register(g) 89 atexit.register(h) 90 91 self.assertRaises(ZeroDivisionError, atexit._run_exitfuncs) 92 stderr = self.stream.getvalue() 93 self.assertEqual(stderr.count("ZeroDivisionError"), 3) 94 self.assertIn("# one", stderr) 95 self.assertIn("# two", stderr) 96 self.assertIn("# three", stderr) 97 98 def test_stress(self): 99 a = [0] 100 def inc(): 101 a[0] += 1 102 103 for i in range(128): 104 atexit.register(inc) 105 atexit._run_exitfuncs() 106 107 self.assertEqual(a[0], 128) 108 109 def test_clear(self): 110 a = [0] 111 def inc(): 112 a[0] += 1 113 114 atexit.register(inc) 115 atexit._clear() 116 atexit._run_exitfuncs() 117 118 self.assertEqual(a[0], 0) 119 120 def test_unregister(self): 121 a = [0] 122 def inc(): 123 a[0] += 1 124 def dec(): 125 a[0] -= 1 126 127 for i in range(4): 128 atexit.register(inc) 129 atexit.register(dec) 130 atexit.unregister(inc) 131 atexit._run_exitfuncs() 132 133 self.assertEqual(a[0], -1) 134 135 def test_bound_methods(self): 136 l = [] 137 atexit.register(l.append, 5) 138 atexit._run_exitfuncs() 139 self.assertEqual(l, [5]) 140 141 atexit.unregister(l.append) 142 atexit._run_exitfuncs() 143 self.assertEqual(l, [5]) 144 145 146class SubinterpreterTest(unittest.TestCase): 147 148 def test_callbacks_leak(self): 149 # This test shows a leak in refleak mode if atexit doesn't 150 # take care to free callbacks in its per-subinterpreter module 151 # state. 152 n = atexit._ncallbacks() 153 code = r"""if 1: 154 import atexit 155 def f(): 156 pass 157 atexit.register(f) 158 del atexit 159 """ 160 ret = support.run_in_subinterp(code) 161 self.assertEqual(ret, 0) 162 self.assertEqual(atexit._ncallbacks(), n) 163 164 def test_callbacks_leak_refcycle(self): 165 # Similar to the above, but with a refcycle through the atexit 166 # module. 167 n = atexit._ncallbacks() 168 code = r"""if 1: 169 import atexit 170 def f(): 171 pass 172 atexit.register(f) 173 atexit.__atexit = atexit 174 """ 175 ret = support.run_in_subinterp(code) 176 self.assertEqual(ret, 0) 177 self.assertEqual(atexit._ncallbacks(), n) 178 179 180if __name__ == "__main__": 181 unittest.main() 182