pytest: add file removal helpers for TestCaseInTempDir

In several places we end a test by deleting a number of files and
directories, but we do it rather haphazardly with unintentionally
differing error handling. For example, in some tests we currently have
something like:

        try:
            shutil.rmtree(os.path.join(self.tempdir, "a"))
            os.remove(os.path.join(self.tempdir, "b"))
            shutil.rmtree(os.path.join(self.tempdir, "c"))
        except Exception:
            pass

where if, for example, the removal of "b" fails, the removal of "c" will
not be attempted. That will result in the tearDown method raising an
exception, and we're no better off. If the above code is replaced with

        self.rm_files('b')
        self.rm_dirs('a', 'c')

the failure to remove 'b' will cause a test error, *unless* the failure
was due to a FileNotFoundError (a.k.a. an OSError with errno ENOENT),
in which case we ignore it, as was probably the original intention.

If on the other hand, we have

        self.rm_files('b', must_exist=True)
        self.rm_dirs('a', 'c')

then the FileNotFoundError causes a failure (not an error).

We take a little bit of care to stay within self.tempdir, to protect
test authors who accidentally write something like `self.rm_dirs('/')`.

Signed-off-by: Douglas Bagnall <douglas.bagnall@catalyst.net.nz>
Reviewed-by: Andrew Bartlett <abartlet@samba.org>
Reviewed-by: Noel Power <npower@samba.org>

1 files changed, 35 insertions, 0 deletions

diff --git a/python/samba/tests/__init__.py b/python/samba/tests/__init__.py
index 3bb7995052c..e37ceac9bc9 100644
--- a/python/samba/tests/__init__.py
+++ b/python/samba/tests/__init__.py
@@ -37,6 +37,7 @@ import samba.dcerpc.base
 from random import randint
 from random import SystemRandom
 from contextlib import contextmanager
+import shutil
 import string
 try:
     from samba.samdb import SamDB
@@ -295,6 +296,40 @@ class TestCaseInTempDir(TestCase):
                 print("could not remove temporary file: %s" % e,
                       file=sys.stderr)
 
+    def rm_files(self, *files, allow_missing=False, _rm=os.remove):
+        """Remove listed files from the temp directory.
+
+        The files must be true files in the directory itself, not in
+        sub-directories.
+
+        By default a non-existent file will cause a test failure (or
+        error if used outside a test in e.g. tearDown), but if
+        allow_missing is true, the absence will be ignored.
+        """
+        for f in files:
+            path = os.path.join(self.tempdir, f)
+
+            # os.path.join will happily step out of the tempdir,
+            # so let's just check.
+            if os.path.dirname(path) != self.tempdir:
+                raise ValueError("{path} might be outside {self.tempdir}")
+
+            try:
+                _rm(path)
+            except FileNotFoundError as e:
+                if not allow_missing:
+                    raise AssertionError(f"{f} not in {self.tempdir}: {e}")
+
+                print(f"{f} not in {self.tempdir}")
+
+    def rm_dirs(self, *dirs, allow_missing=False):
+        """Remove listed directories from temp directory.
+
+        This works like rm_files, but only removes directories,
+        including their contents.
+        """
+        self.rm_files(*dirs, allow_missing=allow_missing, _rm=shutil.rmtree)
+
 
 def env_loadparm():
     lp = param.LoadParm()