Close #15415: Factor out temp dir helpers to test.support

Patch by Chris Jerdonek

5 files changed, 190 insertions, 46 deletions

diff --git a/Lib/test/script_helper.py b/Lib/test/script_helper.py
index b09f4bf..ab20164 100644
--- a/Lib/test/script_helper.py
+++ b/Lib/test/script_helper.py
@@ -13,7 +13,7 @@ import shutil
 import zipfile
 
 from imp import source_from_cache
-from test.support import make_legacy_pyc, strip_python_stderr
+from test.support import make_legacy_pyc, strip_python_stderr, temp_dir
 
 # Executing the interpreter in a subprocess
 def _assert_python(expected_success, *args, **env_vars):
@@ -77,16 +77,6 @@ def kill_python(p):
     subprocess._cleanup()
     return data
 
-# Script creation utilities
-@contextlib.contextmanager
-def temp_dir():
-    dirname = tempfile.mkdtemp()
-    dirname = os.path.realpath(dirname)
-    try:
-        yield dirname
-    finally:
-        shutil.rmtree(dirname)
-
 def make_script(script_dir, script_basename, source):
     script_filename = script_basename+os.extsep+'py'
     script_name = os.path.join(script_dir, script_filename)
diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py
index 3ebc49d..59be326 100644
--- a/Lib/test/support/__init__.py
+++ b/Lib/test/support/__init__.py
@@ -738,45 +738,85 @@ else:
 SAVEDCWD = os.getcwd()
 
 @contextlib.contextmanager
-def temp_cwd(name='tempcwd', quiet=False, path=None):
-    """
-    Context manager that temporarily changes the CWD.
+def temp_dir(path=None, quiet=False):
+    """Return a context manager that creates a temporary directory.
+
+    Arguments:
+
+      path: the directory to create temporarily.  If omitted or None,
+        defaults to creating a temporary directory using tempfile.mkdtemp.
 
-    An existing path may be provided as *path*, in which case this
-    function makes no changes to the file system.
+      quiet: if False (the default), the context manager raises an exception
+        on error.  Otherwise, if the path is specified and cannot be
+        created, only a warning is issued.
 
-    Otherwise, the new CWD is created in the current directory and it's
-    named *name*. If *quiet* is False (default) and it's not possible to
-    create or change the CWD, an error is raised.  If it's True, only a
-    warning is raised and the original CWD is used.
     """
-    saved_dir = os.getcwd()
-    is_temporary = False
+    dir_created = False
     if path is None:
-        path = name
+        path = tempfile.mkdtemp()
+        dir_created = True
+        path = os.path.realpath(path)
+    else:
         try:
-            os.mkdir(name)
-            is_temporary = True
+            os.mkdir(path)
+            dir_created = True
         except OSError:
             if not quiet:
                 raise
-            warnings.warn('tests may fail, unable to create temp CWD ' + name,
+            warnings.warn('tests may fail, unable to create temp dir: ' + path,
                           RuntimeWarning, stacklevel=3)
     try:
+        yield path
+    finally:
+        if dir_created:
+            shutil.rmtree(path)
+
+@contextlib.contextmanager
+def change_cwd(path, quiet=False):
+    """Return a context manager that changes the current working directory.
+
+    Arguments:
+
+      path: the directory to use as the temporary current working directory.
+
+      quiet: if False (the default), the context manager raises an exception
+        on error.  Otherwise, it issues only a warning and keeps the current
+        working directory the same.
+
+    """
+    saved_dir = os.getcwd()
+    try:
         os.chdir(path)
     except OSError:
         if not quiet:
             raise
-        warnings.warn('tests may fail, unable to change the CWD to ' + path,
+        warnings.warn('tests may fail, unable to change CWD to: ' + path,
                       RuntimeWarning, stacklevel=3)
     try:
         yield os.getcwd()
     finally:
         os.chdir(saved_dir)
-        if is_temporary:
-            rmtree(name)
 
 
+@contextlib.contextmanager
+def temp_cwd(name='tempcwd', quiet=False):
+    """
+    Context manager that temporarily creates and changes the CWD.
+
+    The function temporarily changes the current working directory
+    after creating a temporary directory in the current directory with
+    name *name*.  If *name* is None, the temporary directory is
+    created using tempfile.mkdtemp.
+
+    If *quiet* is False (default) and it is not possible to
+    create or change the CWD, an error is raised.  If *quiet* is True,
+    only a warning is raised and the original CWD is used.
+
+    """
+    with temp_dir(path=name, quiet=quiet) as temp_path:
+        with change_cwd(temp_path, quiet=quiet) as cwd_dir:
+            yield cwd_dir
+
 if hasattr(os, "umask"):
     @contextlib.contextmanager
     def temp_umask(umask):
diff --git a/Lib/test/test_cmd_line_script.py b/Lib/test/test_cmd_line_script.py
index 6051e18..b4269b9 100644
--- a/Lib/test/test_cmd_line_script.py
+++ b/Lib/test/test_cmd_line_script.py
@@ -290,7 +290,7 @@ class CmdLineTest(unittest.TestCase):
         # Make sure package __init__ modules see "-m" in sys.argv0 while
         # searching for the module to execute
         with temp_dir() as script_dir:
-            with support.temp_cwd(path=script_dir):
+            with support.change_cwd(path=script_dir):
                 pkg_dir = os.path.join(script_dir, 'test_pkg')
                 make_pkg(pkg_dir, "import sys; print('init_argv0==%r' % sys.argv[0])")
                 script_name = _make_test_script(pkg_dir, 'script')
@@ -307,7 +307,7 @@ class CmdLineTest(unittest.TestCase):
         # Make sure a "-c" file in the current directory
         # does not alter the value of sys.path[0]
         with temp_dir() as script_dir:
-            with support.temp_cwd(path=script_dir):
+            with support.change_cwd(path=script_dir):
                 with open("-c", "w") as f:
                     f.write("data")
                     rc, out, err = assert_python_ok('-c',
@@ -322,7 +322,7 @@ class CmdLineTest(unittest.TestCase):
         # does not alter the value of sys.path[0]
         with temp_dir() as script_dir:
             script_name = _make_test_script(script_dir, 'other')
-            with support.temp_cwd(path=script_dir):
+            with support.change_cwd(path=script_dir):
                 with open("-m", "w") as f:
                     f.write("data")
                     rc, out, err = assert_python_ok('-m', 'other', *example_args)
@@ -335,7 +335,7 @@ class CmdLineTest(unittest.TestCase):
         # and results in an error that the return code to the
         # shell is '1'
         with temp_dir() as script_dir:
-            with support.temp_cwd(path=script_dir):
+            with support.change_cwd(path=script_dir):
                 pkg_dir = os.path.join(script_dir, 'test_pkg')
                 make_pkg(pkg_dir)
                 script_name = _make_test_script(pkg_dir, 'other',
diff --git a/Lib/test/test_shutil.py b/Lib/test/test_shutil.py
index 6794a93..acaffdd 100644
--- a/Lib/test/test_shutil.py
+++ b/Lib/test/test_shutil.py
@@ -1307,18 +1307,18 @@ class TestWhich(unittest.TestCase):
         # that exists, it should be returned.
         base_dir, tail_dir = os.path.split(self.dir)
         relpath = os.path.join(tail_dir, self.file)
-        with support.temp_cwd(path=base_dir):
+        with support.change_cwd(path=base_dir):
             rv = shutil.which(relpath, path=self.temp_dir)
             self.assertEqual(rv, relpath)
         # But it shouldn't be searched in PATH directories (issue #16957).
-        with support.temp_cwd(path=self.dir):
+        with support.change_cwd(path=self.dir):
             rv = shutil.which(relpath, path=base_dir)
             self.assertIsNone(rv)
 
     def test_cwd(self):
         # Issue #16957
         base_dir = os.path.dirname(self.dir)
-        with support.temp_cwd(path=self.dir):
+        with support.change_cwd(path=self.dir):
             rv = shutil.which(self.file, path=base_dir)
             if sys.platform == "win32":
                 # Windows: current directory implicitly on PATH
@@ -1339,7 +1339,7 @@ class TestWhich(unittest.TestCase):
 
     def test_relative_path(self):
         base_dir, tail_dir = os.path.split(self.dir)
-        with support.temp_cwd(path=base_dir):
+        with support.change_cwd(path=base_dir):
             rv = shutil.which(self.file, path=tail_dir)
             self.assertEqual(rv, os.path.join(tail_dir, self.file))
 
@@ -1364,7 +1364,7 @@ class TestWhich(unittest.TestCase):
 
     def test_empty_path(self):
         base_dir = os.path.dirname(self.dir)
-        with support.temp_cwd(path=self.dir), \
+        with support.change_cwd(path=self.dir), \
              support.EnvironmentVarGuard() as env:
             env['PATH'] = self.dir
             rv = shutil.which(self.file, path='')
diff --git a/Lib/test/test_support.py b/Lib/test/test_support.py
index 340b8da..4edb1a8 100644
--- a/Lib/test/test_support.py
+++ b/Lib/test/test_support.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
 import importlib
+import shutil
 import sys
 import os
 import unittest
@@ -88,6 +89,118 @@ class TestSupport(unittest.TestCase):
         s.listen(1)
         s.close()
 
+    # Tests for temp_dir()
+
+    def test_temp_dir(self):
+        """Test that temp_dir() creates and destroys its directory."""
+        parent_dir = tempfile.mkdtemp()
+        parent_dir = os.path.realpath(parent_dir)
+
+        try:
+            path = os.path.join(parent_dir, 'temp')
+            self.assertFalse(os.path.isdir(path))
+            with support.temp_dir(path) as temp_path:
+                self.assertEqual(temp_path, path)
+                self.assertTrue(os.path.isdir(path))
+            self.assertFalse(os.path.isdir(path))
+        finally:
+            shutil.rmtree(parent_dir)
+
+    def test_temp_dir__path_none(self):
+        """Test passing no path."""
+        with support.temp_dir() as temp_path:
+            self.assertTrue(os.path.isdir(temp_path))
+        self.assertFalse(os.path.isdir(temp_path))
+
+    def test_temp_dir__existing_dir__quiet_default(self):
+        """Test passing a directory that already exists."""
+        def call_temp_dir(path):
+            with support.temp_dir(path) as temp_path:
+                raise Exception("should not get here")
+
+        path = tempfile.mkdtemp()
+        path = os.path.realpath(path)
+        try:
+            self.assertTrue(os.path.isdir(path))
+            self.assertRaises(FileExistsError, call_temp_dir, path)
+            # Make sure temp_dir did not delete the original directory.
+            self.assertTrue(os.path.isdir(path))
+        finally:
+            shutil.rmtree(path)
+
+    def test_temp_dir__existing_dir__quiet_true(self):
+        """Test passing a directory that already exists with quiet=True."""
+        path = tempfile.mkdtemp()
+        path = os.path.realpath(path)
+
+        try:
+            with support.check_warnings() as recorder:
+                with support.temp_dir(path, quiet=True) as temp_path:
+                    self.assertEqual(path, temp_path)
+                warnings = [str(w.message) for w in recorder.warnings]
+            # Make sure temp_dir did not delete the original directory.
+            self.assertTrue(os.path.isdir(path))
+        finally:
+            shutil.rmtree(path)
+
+        expected = ['tests may fail, unable to create temp dir: ' + path]
+        self.assertEqual(warnings, expected)
+
+    # Tests for change_cwd()
+
+    def test_change_cwd(self):
+        original_cwd = os.getcwd()
+
+        with support.temp_dir() as temp_path:
+            with support.change_cwd(temp_path) as new_cwd:
+                self.assertEqual(new_cwd, temp_path)
+                self.assertEqual(os.getcwd(), new_cwd)
+
+        self.assertEqual(os.getcwd(), original_cwd)
+
+    def test_change_cwd__non_existent_dir(self):
+        """Test passing a non-existent directory."""
+        original_cwd = os.getcwd()
+
+        def call_change_cwd(path):
+            with support.change_cwd(path) as new_cwd:
+                raise Exception("should not get here")
+
+        with support.temp_dir() as parent_dir:
+            non_existent_dir = os.path.join(parent_dir, 'does_not_exist')
+            self.assertRaises(FileNotFoundError, call_change_cwd,
+                              non_existent_dir)
+
+        self.assertEqual(os.getcwd(), original_cwd)
+
+    def test_change_cwd__non_existent_dir__quiet_true(self):
+        """Test passing a non-existent directory with quiet=True."""
+        original_cwd = os.getcwd()
+
+        with support.temp_dir() as parent_dir:
+            bad_dir = os.path.join(parent_dir, 'does_not_exist')
+            with support.check_warnings() as recorder:
+                with support.change_cwd(bad_dir, quiet=True) as new_cwd:
+                    self.assertEqual(new_cwd, original_cwd)
+                    self.assertEqual(os.getcwd(), new_cwd)
+                warnings = [str(w.message) for w in recorder.warnings]
+
+        expected = ['tests may fail, unable to change CWD to: ' + bad_dir]
+        self.assertEqual(warnings, expected)
+
+    # Tests for change_cwd()
+
+    def test_change_cwd__chdir_warning(self):
+        """Check the warning message when os.chdir() fails."""
+        path = TESTFN + '_does_not_exist'
+        with support.check_warnings() as recorder:
+            with support.change_cwd(path=path, quiet=True):
+                pass
+            messages = [str(w.message) for w in recorder.warnings]
+        self.assertEqual(messages, ['tests may fail, unable to change CWD to: ' + path])
+
+    # Tests for temp_cwd()
+
     def test_temp_cwd(self):
         here = os.getcwd()
         with support.temp_cwd(name=TESTFN):
@@ -95,14 +208,15 @@ class TestSupport(unittest.TestCase):
         self.assertFalse(os.path.exists(TESTFN))
         self.assertTrue(os.path.basename(os.getcwd()), here)
 
-    def test_temp_cwd__chdir_warning(self):
-        """Check the warning message when os.chdir() fails."""
-        path = TESTFN + '_does_not_exist'
-        with support.check_warnings() as recorder:
-            with support.temp_cwd(path=path, quiet=True):
-                pass
-            messages = [str(w.message) for w in recorder.warnings]
-        self.assertEqual(messages, ['tests may fail, unable to change the CWD to ' + path])
+
+    def test_temp_cwd__name_none(self):
+        """Test passing None to temp_cwd()."""
+        original_cwd = os.getcwd()
+        with support.temp_cwd(name=None) as new_cwd:
+            self.assertNotEqual(new_cwd, original_cwd)
+            self.assertTrue(os.path.isdir(new_cwd))
+            self.assertEqual(os.getcwd(), new_cwd)
+        self.assertEqual(os.getcwd(), original_cwd)
 
     def test_sortdict(self):
         self.assertEqual(support.sortdict({3:3, 2:2, 1:1}), "{1: 1, 2: 2, 3: 3}")