[testing] a new TestCasePlus subclass + get_auto_remove_tmp_dir() (#6494)

* [testing] switch to a new TestCasePlus + get_auto_remove_tmp_dir() for auto-removal of tmp dirs

* respect after=True for tempfile, simplify code

* comments

* comment fix

* put `before` last in args, so can make debug even faster

src/transformers/testing_utils.py

@@ -1,9 +1,12 @@
 import os
 import re
+import shutil
 import sys
+import tempfile
 import unittest
 from distutils.util import strtobool
 from io import StringIO
+from pathlib import Path
 
 from .file_utils import _tf_available, _torch_available, _torch_tpu_available
 
@@ -255,3 +258,92 @@ class CaptureStderr(CaptureStd):
 
     def __init__(self):
         super().__init__(out=False)
+
+
+class TestCasePlus(unittest.TestCase):
+    """This class extends `unittest.TestCase` with additional features.
+
+    Feature 1: Flexible auto-removable temp dirs which are guaranteed to get
+    removed at the end of test.
+
+    In all the following scenarios the temp dir will be auto-removed at the end
+    of test, unless `after=False`.
+
+    # 1. create a unique temp dir, `tmp_dir` will contain the path to the created temp dir
+    def test_whatever(self):
+        tmp_dir = self.get_auto_remove_tmp_dir()
+
+    # 2. create a temp dir of my choice and delete it at the end - useful for debug when you want to
+    # monitor a specific directory
+    def test_whatever(self):
+        tmp_dir = self.get_auto_remove_tmp_dir(tmp_dir="./tmp/run/test")
+
+    # 3. create a temp dir of my choice and do not delete it at the end - useful for when you want
+    # to look at the temp results
+    def test_whatever(self):
+        tmp_dir = self.get_auto_remove_tmp_dir(tmp_dir="./tmp/run/test", after=False)
+
+    # 4. create a temp dir of my choice and ensure to delete it right away - useful for when you
+    # disabled deletion in the previous test run and want to make sure the that tmp dir is empty
+    # before the new test is run
+    def test_whatever(self):
+        tmp_dir = self.get_auto_remove_tmp_dir(tmp_dir="./tmp/run/test", before=True)
+
+    Note 1: In order to run the equivalent of `rm -r` safely, only subdirs of the
+    project repository checkout are allowed if an explicit `tmp_dir` is used, so
+    that by mistake no `/tmp` or similar important part of the filesystem will
+    get nuked. i.e. please always pass paths that start with `./`
+
+    Note 2: Each test can register multiple temp dirs and they all will get
+    auto-removed, unless requested otherwise.
+
+    """
+
+    def setUp(self):
+        self.teardown_tmp_dirs = []
+
+    def get_auto_remove_tmp_dir(self, tmp_dir=None, after=True, before=False):
+        """
+        Args:
+            tmp_dir (:obj:`string`, `optional`, defaults to :obj:`None`):
+                use this path, if None a unique path will be assigned
+            before (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                if `True` and tmp dir already exists make sure to empty it right away
+            after (:obj:`bool`, `optional`, defaults to :obj:`True`):
+                delete the tmp dir at the end of the test
+
+        Returns:
+            tmp_dir(:obj:`string`):
+                either the same value as passed via `tmp_dir` or the path to the auto-created tmp dir
+        """
+        if tmp_dir is not None:
+            # using provided path
+            path = Path(tmp_dir).resolve()
+
+            # to avoid nuking parts of the filesystem, only relative paths are allowed
+            if not tmp_dir.startswith("./"):
+                raise ValueError(
+                    f"`tmp_dir` can only be a relative path, i.e. `./some/path`, but received `{tmp_dir}`"
+                )
+
+            # ensure the dir is empty to start with
+            if before is True and path.exists():
+                shutil.rmtree(tmp_dir, ignore_errors=True)
+
+            path.mkdir(parents=True, exist_ok=True)
+
+        else:
+            # using unique tmp dir (always empty, regardless of `before`)
+            tmp_dir = tempfile.mkdtemp()
+
+        if after is True:
+            # register for deletion
+            self.teardown_tmp_dirs.append(tmp_dir)
+
+        return tmp_dir
+
+    def tearDown(self):
+        # remove registered temp dirs
+        for path in self.teardown_tmp_dirs:
+            shutil.rmtree(path, ignore_errors=True)
+        self.teardown_tmp_dirs = []