From b94d7fd4efa877d649cea9c8125c8869ffe0c32d Mon Sep 17 00:00:00 2001
From: Zachary Ware <zachary.ware@gmail.com>
Date: Sat, 1 Apr 2017 00:18:23 -0500
Subject: [PATCH] bpo-29763: Use unittest cleanup in test_site (GH-841)

---
 Lib/test/test_site.py | 79 +++++++++++++++++++------------------------
 1 file changed, 34 insertions(+), 45 deletions(-)

diff --git a/Lib/test/test_site.py b/Lib/test/test_site.py
index 342ec9e43d5..78d59e290da 100644
--- a/Lib/test/test_site.py
+++ b/Lib/test/test_site.py
@@ -17,6 +17,7 @@ import urllib.error
 import shutil
 import subprocess
 import sysconfig
+import tempfile
 from copy import copy
 
 # These tests are not particularly useful if Python was invoked with -S.
@@ -489,29 +490,25 @@ class StartupImportTests(unittest.TestCase):
             'import site, sys; site.enablerlcompleter(); sys.exit(hasattr(sys, "__interactivehook__"))']).wait()
         self.assertTrue(r, "'__interactivehook__' not added by enablerlcompleter()")
 
-    @classmethod
+
+@unittest.skipUnless(sys.platform == 'win32', "only supported on Windows")
+class _pthFileTests(unittest.TestCase):
+
     def _create_underpth_exe(self, lines):
-        exe_file = os.path.join(os.getenv('TEMP'), os.path.split(sys.executable)[1])
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.addCleanup(self.temp_dir.cleanup)
+        exe_file = os.path.join(
+            self.temp_dir.name,
+            os.path.split(sys.executable)[1],
+        )
         shutil.copy(sys.executable, exe_file)
 
         _pth_file = os.path.splitext(exe_file)[0] + '._pth'
-        try:
-            with open(_pth_file, 'w') as f:
-                for line in lines:
-                    print(line, file=f)
-            return exe_file
-        except:
-            test.support.unlink(_pth_file)
-            test.support.unlink(exe_file)
-            raise
+        with open(_pth_file, 'w') as f:
+            for line in lines:
+                print(line, file=f)
+        return exe_file
 
-    @classmethod
-    def _cleanup_underpth_exe(self, exe_file):
-        _pth_file = os.path.splitext(exe_file)[0] + '._pth'
-        test.support.unlink(_pth_file)
-        test.support.unlink(exe_file)
-
-    @classmethod
     def _calc_sys_path_for_underpth_nosite(self, sys_prefix, lines):
         sys_path = []
         for line in lines:
@@ -521,7 +518,6 @@ class StartupImportTests(unittest.TestCase):
             sys_path.append(abs_path)
         return sys_path
 
-    @unittest.skipUnless(sys.platform == 'win32', "only supported on Windows")
     def test_underpth_nosite_file(self):
         libpath = os.path.dirname(os.path.dirname(encodings.__file__))
         exe_prefix = os.path.dirname(sys.executable)
@@ -536,20 +532,16 @@ class StartupImportTests(unittest.TestCase):
             os.path.dirname(exe_file),
             pth_lines)
 
-        try:
-            env = os.environ.copy()
-            env['PYTHONPATH'] = 'from-env'
-            env['PATH'] = '{};{}'.format(exe_prefix, os.getenv('PATH'))
-            rc = subprocess.call([exe_file, '-c',
-                'import sys; sys.exit(sys.flags.no_site and '
-                'len(sys.path) > 200 and '
-                'sys.path == %r)' % sys_path,
-                ], env=env)
-        finally:
-            self._cleanup_underpth_exe(exe_file)
+        env = os.environ.copy()
+        env['PYTHONPATH'] = 'from-env'
+        env['PATH'] = '{};{}'.format(exe_prefix, os.getenv('PATH'))
+        rc = subprocess.call([exe_file, '-c',
+            'import sys; sys.exit(sys.flags.no_site and '
+            'len(sys.path) > 200 and '
+            'sys.path == %r)' % sys_path,
+            ], env=env)
         self.assertTrue(rc, "sys.path is incorrect")
 
-    @unittest.skipUnless(sys.platform == 'win32', "only supported on Windows")
     def test_underpth_file(self):
         libpath = os.path.dirname(os.path.dirname(encodings.__file__))
         exe_prefix = os.path.dirname(sys.executable)
@@ -561,20 +553,17 @@ class StartupImportTests(unittest.TestCase):
             'import site'
         ])
         sys_prefix = os.path.dirname(exe_file)
-        try:
-            env = os.environ.copy()
-            env['PYTHONPATH'] = 'from-env'
-            env['PATH'] = '{};{}'.format(exe_prefix, os.getenv('PATH'))
-            rc = subprocess.call([exe_file, '-c',
-                'import sys; sys.exit(not sys.flags.no_site and '
-                '%r in sys.path and %r in sys.path and %r not in sys.path and '
-                'all("\\r" not in p and "\\n" not in p for p in sys.path))' % (
-                    os.path.join(sys_prefix, 'fake-path-name'),
-                    libpath,
-                    os.path.join(sys_prefix, 'from-env'),
-                )], env=env)
-        finally:
-            self._cleanup_underpth_exe(exe_file)
+        env = os.environ.copy()
+        env['PYTHONPATH'] = 'from-env'
+        env['PATH'] = '{};{}'.format(exe_prefix, os.getenv('PATH'))
+        rc = subprocess.call([exe_file, '-c',
+            'import sys; sys.exit(not sys.flags.no_site and '
+            '%r in sys.path and %r in sys.path and %r not in sys.path and '
+            'all("\\r" not in p and "\\n" not in p for p in sys.path))' % (
+                os.path.join(sys_prefix, 'fake-path-name'),
+                libpath,
+                os.path.join(sys_prefix, 'from-env'),
+            )], env=env)
         self.assertTrue(rc, "sys.path is incorrect")