Issue #18945: Add tests for tempfile name collision handling.

Patch by Vlad Shcherbina

1 files changed, 65 insertions, 22 deletions

diff --git a/Lib/test/test_tempfile.py b/Lib/test/test_tempfile.py
index 0f514c21d1..1ee36e223b 100644
--- a/Lib/test/test_tempfile.py
+++ b/Lib/test/test_tempfile.py
@@ -7,6 +7,7 @@ import signal
 import sys
 import re
 import warnings
+import contextlib
 
 import unittest
 from test import support
@@ -255,6 +256,22 @@ class TestGetCandidateNames(BaseTestCase):
         self.assertTrue(a is b)
 
 
+@contextlib.contextmanager
+def _inside_empty_temp_dir():
+    dir = tempfile.mkdtemp()
+    try:
+        with support.swap_attr(tempfile, 'tempdir', dir):
+            yield
+    finally:
+        support.rmtree(dir)
+
+
+def _mock_candidate_names(*names):
+    return support.swap_attr(tempfile,
+                             '_get_candidate_names',
+                             lambda: iter(names))
+
+
 class TestMkstempInner(BaseTestCase):
     """Test the internal function _mkstemp_inner."""
 
@@ -373,31 +390,36 @@ class TestMkstempInner(BaseTestCase):
         os.lseek(f.fd, 0, os.SEEK_SET)
         self.assertEqual(os.read(f.fd, 20), b"blat")
 
+    def default_mkstemp_inner(self):
+        return tempfile._mkstemp_inner(tempfile.gettempdir(),
+                                       tempfile.template,
+                                       '',
+                                       tempfile._bin_openflags)
+
+    def test_collision_with_existing_file(self):
+        # _mkstemp_inner tries another name when a file with
+        # the chosen name already exists
+        with _inside_empty_temp_dir(), \
+             _mock_candidate_names('aaa', 'aaa', 'bbb'):
+            (fd1, name1) = self.default_mkstemp_inner()
+            os.close(fd1)
+            self.assertTrue(name1.endswith('aaa'))
+
+            (fd2, name2) = self.default_mkstemp_inner()
+            os.close(fd2)
+            self.assertTrue(name2.endswith('bbb'))
+
     def test_collision_with_existing_directory(self):
         # _mkstemp_inner tries another name when a directory with
         # the chosen name already exists
-        container_dir = tempfile.mkdtemp()
-        try:
-            def mock_get_candidate_names():
-                return iter(['aaa', 'aaa', 'bbb'])
-            with support.swap_attr(tempfile,
-                                   '_get_candidate_names',
-                                   mock_get_candidate_names):
-                dir = tempfile.mkdtemp(dir=container_dir)
-                self.assertTrue(dir.endswith('aaa'))
-
-                flags = tempfile._bin_openflags
-                (fd, name) = tempfile._mkstemp_inner(container_dir,
-                                                     tempfile.template,
-                                                     '',
-                                                     flags)
-                try:
-                    self.assertTrue(name.endswith('bbb'))
-                finally:
-                    os.close(fd)
-                    os.unlink(name)
-        finally:
-            support.rmtree(container_dir)
+        with _inside_empty_temp_dir(), \
+             _mock_candidate_names('aaa', 'aaa', 'bbb'):
+            dir = tempfile.mkdtemp()
+            self.assertTrue(dir.endswith('aaa'))
+
+            (fd, name) = self.default_mkstemp_inner()
+            os.close(fd)
+            self.assertTrue(name.endswith('bbb'))
 
 
 class TestGetTempPrefix(BaseTestCase):
@@ -554,6 +576,27 @@ class TestMkdtemp(BaseTestCase):
         finally:
             os.rmdir(dir)
 
+    def test_collision_with_existing_file(self):
+        # mkdtemp tries another name when a file with
+        # the chosen name already exists
+        with _inside_empty_temp_dir(), \
+             _mock_candidate_names('aaa', 'aaa', 'bbb'):
+            file = tempfile.NamedTemporaryFile(delete=False)
+            file.close()
+            self.assertTrue(file.name.endswith('aaa'))
+            dir = tempfile.mkdtemp()
+            self.assertTrue(dir.endswith('bbb'))
+
+    def test_collision_with_existing_directory(self):
+        # mkdtemp tries another name when a directory with
+        # the chosen name already exists
+        with _inside_empty_temp_dir(), \
+             _mock_candidate_names('aaa', 'aaa', 'bbb'):
+            dir1 = tempfile.mkdtemp()
+            self.assertTrue(dir1.endswith('aaa'))
+            dir2 = tempfile.mkdtemp()
+            self.assertTrue(dir2.endswith('bbb'))
+
 
 class TestMktemp(BaseTestCase):
     """Test mktemp()."""