1  import atexit
       2  import os
       3  import textwrap
       4  import unittest
       5  from test import support
       6  from test.support import script_helper
       7  
       8  
       9  class ESC[4;38;5;81mGeneralTest(ESC[4;38;5;149munittestESC[4;38;5;149m.ESC[4;38;5;149mTestCase):
      10      def test_general(self):
      11          # Run _test_atexit.py in a subprocess since it calls atexit._clear()
      12          script = support.findfile("_test_atexit.py")
      13          script_helper.run_test_script(script)
      14  
      15  class ESC[4;38;5;81mFunctionalTest(ESC[4;38;5;149munittestESC[4;38;5;149m.ESC[4;38;5;149mTestCase):
      16      def test_shutdown(self):
      17          # Actually test the shutdown mechanism in a subprocess
      18          code = textwrap.dedent("""
      19              import atexit
      20  
      21              def f(msg):
      22                  print(msg)
      23  
      24              atexit.register(f, "one")
      25              atexit.register(f, "two")
      26          """)
      27          res = script_helper.assert_python_ok("-c", code)
      28          self.assertEqual(res.out.decode().splitlines(), ["two", "one"])
      29          self.assertFalse(res.err)
      30  
      31      def test_atexit_instances(self):
      32          # bpo-42639: It is safe to have more than one atexit instance.
      33          code = textwrap.dedent("""
      34              import sys
      35              import atexit as atexit1
      36              del sys.modules['atexit']
      37              import atexit as atexit2
      38              del sys.modules['atexit']
      39  
      40              assert atexit2 is not atexit1
      41  
      42              atexit1.register(print, "atexit1")
      43              atexit2.register(print, "atexit2")
      44          """)
      45          res = script_helper.assert_python_ok("-c", code)
      46          self.assertEqual(res.out.decode().splitlines(), ["atexit2", "atexit1"])
      47          self.assertFalse(res.err)
      48  
      49  
      50  @support.cpython_only
      51  class ESC[4;38;5;81mSubinterpreterTest(ESC[4;38;5;149munittestESC[4;38;5;149m.ESC[4;38;5;149mTestCase):
      52  
      53      def test_callbacks_leak(self):
      54          # This test shows a leak in refleak mode if atexit doesn't
      55          # take care to free callbacks in its per-subinterpreter module
      56          # state.
      57          n = atexit._ncallbacks()
      58          code = textwrap.dedent(r"""
      59              import atexit
      60              def f():
      61                  pass
      62              atexit.register(f)
      63              del atexit
      64          """)
      65          ret = support.run_in_subinterp(code)
      66          self.assertEqual(ret, 0)
      67          self.assertEqual(atexit._ncallbacks(), n)
      68  
      69      def test_callbacks_leak_refcycle(self):
      70          # Similar to the above, but with a refcycle through the atexit
      71          # module.
      72          n = atexit._ncallbacks()
      73          code = textwrap.dedent(r"""
      74              import atexit
      75              def f():
      76                  pass
      77              atexit.register(f)
      78              atexit.__atexit = atexit
      79          """)
      80          ret = support.run_in_subinterp(code)
      81          self.assertEqual(ret, 0)
      82          self.assertEqual(atexit._ncallbacks(), n)
      83  
      84      @unittest.skipUnless(hasattr(os, "pipe"), "requires os.pipe()")
      85      def test_callback_on_subinterpreter_teardown(self):
      86          # This tests if a callback is called on
      87          # subinterpreter teardown.
      88          expected = b"The test has passed!"
      89          r, w = os.pipe()
      90  
      91          code = textwrap.dedent(r"""
      92              import os
      93              import atexit
      94              def callback():
      95                  os.write({:d}, b"The test has passed!")
      96              atexit.register(callback)
      97          """.format(w))
      98          ret = support.run_in_subinterp(code)
      99          os.close(w)
     100          self.assertEqual(os.read(r, len(expected)), expected)
     101          os.close(r)
     102  
     103  
     104  if __name__ == "__main__":
     105      unittest.main()