4 files changed, 20 insertions, 40 deletions

diff --git a/tests/test_git_superproject.py b/tests/test_git_superproject.py
index 1e7b1201..603694d2 100644
--- a/tests/test_git_superproject.py
+++ b/tests/test_git_superproject.py
@@ -24,7 +24,6 @@ from unittest import mock
 import git_superproject
 import git_trace2_event_log
 import manifest_xml
-import platform_utils
 from test_manifest_xml import sort_attributes
 
 
@@ -38,7 +37,8 @@ class SuperprojectTestCase(unittest.TestCase):
 
   def setUp(self):
     """Set up superproject every time."""
-    self.tempdir = tempfile.mkdtemp(prefix='repo_tests')
+    self.tempdirobj = tempfile.TemporaryDirectory(prefix='repo_tests')
+    self.tempdir = self.tempdirobj.name
     self.repodir = os.path.join(self.tempdir, '.repo')
     self.manifest_file = os.path.join(
         self.repodir, manifest_xml.MANIFEST_FILE_NAME)
@@ -75,7 +75,7 @@ class SuperprojectTestCase(unittest.TestCase):
 
   def tearDown(self):
     """Tear down superproject every time."""
-    platform_utils.rmtree(self.tempdir)
+    self.tempdirobj.cleanup()
 
   def getXmlManifest(self, data):
     """Helper to initialize a manifest for testing."""
diff --git a/tests/test_manifest_xml.py b/tests/test_manifest_xml.py
index ede41547..85c20733 100644
--- a/tests/test_manifest_xml.py
+++ b/tests/test_manifest_xml.py
@@ -17,7 +17,6 @@
 import os
 import platform
 import re
-import shutil
 import tempfile
 import unittest
 import xml.dom.minidom
@@ -92,7 +91,8 @@ class ManifestParseTestCase(unittest.TestCase):
   """TestCase for parsing manifests."""
 
   def setUp(self):
-    self.tempdir = tempfile.mkdtemp(prefix='repo_tests')
+    self.tempdirobj = tempfile.TemporaryDirectory(prefix='repo_tests')
+    self.tempdir = self.tempdirobj.name
     self.repodir = os.path.join(self.tempdir, '.repo')
     self.manifest_dir = os.path.join(self.repodir, 'manifests')
     self.manifest_file = os.path.join(
@@ -111,7 +111,7 @@ class ManifestParseTestCase(unittest.TestCase):
 """)
 
   def tearDown(self):
-    shutil.rmtree(self.tempdir, ignore_errors=True)
+    self.tempdirobj.cleanup()
 
   def getXmlManifest(self, data):
     """Helper to initialize a manifest for testing."""
diff --git a/tests/test_project.py b/tests/test_project.py
index 8829637e..acd44ccc 100644
--- a/tests/test_project.py
+++ b/tests/test_project.py
@@ -17,7 +17,6 @@
 import contextlib
 import os
 from pathlib import Path
-import shutil
 import subprocess
 import tempfile
 import unittest
@@ -32,11 +31,7 @@ import project
 @contextlib.contextmanager
 def TempGitTree():
   """Create a new empty git checkout for testing."""
-  # TODO(vapier): Convert this to tempfile.TemporaryDirectory once we drop
-  # Python 2 support entirely.
-  try:
-    tempdir = tempfile.mkdtemp(prefix='repo-tests')
-
+  with tempfile.TemporaryDirectory(prefix='repo-tests') as tempdir:
     # Tests need to assume, that main is default branch at init,
     # which is not supported in config until 2.28.
     cmd = ['git', 'init']
@@ -50,8 +45,6 @@ def TempGitTree():
       cmd += ['--template', templatedir]
     subprocess.check_call(cmd, cwd=tempdir)
     yield tempdir
-  finally:
-    platform_utils.rmtree(tempdir)
 
 
 class FakeProject(object):
@@ -124,14 +117,15 @@ class CopyLinkTestCase(unittest.TestCase):
   """
 
   def setUp(self):
-    self.tempdir = tempfile.mkdtemp(prefix='repo_tests')
+    self.tempdirobj = tempfile.TemporaryDirectory(prefix='repo_tests')
+    self.tempdir = self.tempdirobj.name
     self.topdir = os.path.join(self.tempdir, 'checkout')
     self.worktree = os.path.join(self.topdir, 'git-project')
     os.makedirs(self.topdir)
     os.makedirs(self.worktree)
 
   def tearDown(self):
-    shutil.rmtree(self.tempdir, ignore_errors=True)
+    self.tempdirobj.cleanup()
 
   @staticmethod
   def touch(path):
diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py
index e9a1f64a..8447bec1 100644
--- a/tests/test_wrapper.py
+++ b/tests/test_wrapper.py
@@ -14,11 +14,9 @@
 
 """Unittests for the wrapper.py module."""
 
-import contextlib
 from io import StringIO
 import os
 import re
-import shutil
 import sys
 import tempfile
 import unittest
@@ -26,22 +24,9 @@ from unittest import mock
 
 import git_command
 import main
-import platform_utils
 import wrapper
 
 
-@contextlib.contextmanager
-def TemporaryDirectory():
-  """Create a new empty git checkout for testing."""
-  # TODO(vapier): Convert this to tempfile.TemporaryDirectory once we drop
-  # Python 2 support entirely.
-  try:
-    tempdir = tempfile.mkdtemp(prefix='repo-tests')
-    yield tempdir
-  finally:
-    platform_utils.rmtree(tempdir)
-
-
 def fixture(*paths):
   """Return a path relative to tests/fixtures.
   """
@@ -336,19 +321,19 @@ class NeedSetupGnuPG(RepoWrapperTestCase):
 
   def test_missing_dir(self):
     """The ~/.repoconfig tree doesn't exist yet."""
-    with TemporaryDirectory() as tempdir:
+    with tempfile.TemporaryDirectory(prefix='repo-tests') as tempdir:
       self.wrapper.home_dot_repo = os.path.join(tempdir, 'foo')
       self.assertTrue(self.wrapper.NeedSetupGnuPG())
 
   def test_missing_keyring(self):
     """The keyring-version file doesn't exist yet."""
-    with TemporaryDirectory() as tempdir:
+    with tempfile.TemporaryDirectory(prefix='repo-tests') as tempdir:
       self.wrapper.home_dot_repo = tempdir
       self.assertTrue(self.wrapper.NeedSetupGnuPG())
 
   def test_empty_keyring(self):
     """The keyring-version file exists, but is empty."""
-    with TemporaryDirectory() as tempdir:
+    with tempfile.TemporaryDirectory(prefix='repo-tests') as tempdir:
       self.wrapper.home_dot_repo = tempdir
       with open(os.path.join(tempdir, 'keyring-version'), 'w'):
         pass
@@ -356,7 +341,7 @@ class NeedSetupGnuPG(RepoWrapperTestCase):
 
   def test_old_keyring(self):
     """The keyring-version file exists, but it's old."""
-    with TemporaryDirectory() as tempdir:
+    with tempfile.TemporaryDirectory(prefix='repo-tests') as tempdir:
       self.wrapper.home_dot_repo = tempdir
       with open(os.path.join(tempdir, 'keyring-version'), 'w') as fp:
         fp.write('1.0\n')
@@ -364,7 +349,7 @@ class NeedSetupGnuPG(RepoWrapperTestCase):
 
   def test_new_keyring(self):
     """The keyring-version file exists, and is up-to-date."""
-    with TemporaryDirectory() as tempdir:
+    with tempfile.TemporaryDirectory(prefix='repo-tests') as tempdir:
       self.wrapper.home_dot_repo = tempdir
       with open(os.path.join(tempdir, 'keyring-version'), 'w') as fp:
         fp.write('1000.0\n')
@@ -376,7 +361,7 @@ class SetupGnuPG(RepoWrapperTestCase):
 
   def test_full(self):
     """Make sure it works completely."""
-    with TemporaryDirectory() as tempdir:
+    with tempfile.TemporaryDirectory(prefix='repo-tests') as tempdir:
       self.wrapper.home_dot_repo = tempdir
       self.wrapper.gpg_dir = os.path.join(self.wrapper.home_dot_repo, 'gnupg')
       self.assertTrue(self.wrapper.SetupGnuPG(True))
@@ -426,7 +411,8 @@ class GitCheckoutTestCase(RepoWrapperTestCase):
   @classmethod
   def setUpClass(cls):
     # Create a repo to operate on, but do it once per-class.
-    cls.GIT_DIR = tempfile.mkdtemp(prefix='repo-rev-tests')
+    cls.tempdirobj = tempfile.TemporaryDirectory(prefix='repo-rev-tests')
+    cls.GIT_DIR = cls.tempdirobj.name
     run_git = wrapper.Wrapper().run_git
 
     remote = os.path.join(cls.GIT_DIR, 'remote')
@@ -455,10 +441,10 @@ class GitCheckoutTestCase(RepoWrapperTestCase):
 
   @classmethod
   def tearDownClass(cls):
-    if not cls.GIT_DIR:
+    if not cls.tempdirobj:
       return
 
-    shutil.rmtree(cls.GIT_DIR)
+    cls.tempdirobj.cleanup()
 
 
 class ResolveRepoRev(GitCheckoutTestCase):