1 files changed, 52 insertions, 14 deletions
diff --git a/libbtrfsutil/python/tests/__init__.py b/libbtrfsutil/python/tests/__init__.py
index 35550e0a..9fd6f6de 100644
--- a/libbtrfsutil/python/tests/__init__.py
+++ b/libbtrfsutil/python/tests/__init__.py
@@ -15,26 +15,60 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with libbtrfsutil.  If not, see <http://www.gnu.org/licenses/>.
 
+import contextlib
 import os
 from pathlib import PurePath
+import pwd
 import subprocess
 import tempfile
 import unittest
 
 
 HAVE_PATH_LIKE = hasattr(PurePath, '__fspath__')
+try:
+    NOBODY_UID = pwd.getpwnam('nobody').pw_uid
+    skipUnlessHaveNobody = lambda func: func
+except KeyError:
+    NOBODY_UID = None
+    skipUnlessHaveNobody = unittest.skip('must have nobody user')
+
+
+@contextlib.contextmanager
+def drop_privs():
+    try:
+        os.seteuid(NOBODY_UID)
+        yield
+    finally:
+        os.seteuid(0)
+
+
+@contextlib.contextmanager
+def regain_privs():
+    uid = os.geteuid()
+    if uid:
+        try:
+            os.seteuid(0)
+            yield
+        finally:
+            os.seteuid(uid)
+    else:
+        yield
 
 
 @unittest.skipIf(os.geteuid() != 0, 'must be run as root')
 class BtrfsTestCase(unittest.TestCase):
-    def setUp(self):
-        self.mountpoint = tempfile.mkdtemp()
+    def __init__(self, *args, **kwds):
+        super().__init__(*args, **kwds)
+        self._mountpoints = []
+
+    def mount_btrfs(self):
+        mountpoint = tempfile.mkdtemp()
         try:
             with tempfile.NamedTemporaryFile(delete=False) as f:
                 os.truncate(f.fileno(), 1024 * 1024 * 1024)
-                self.image = f.name
+                image = f.name
         except Exception as e:
-            os.rmdir(self.mountpoint)
+            os.rmdir(mountpoint)
             raise e
 
         if os.path.exists('../../mkfs.btrfs'):
@@ -42,19 +76,24 @@ class BtrfsTestCase(unittest.TestCase):
         else:
             mkfs = 'mkfs.btrfs'
         try:
-            subprocess.check_call([mkfs, '-q', self.image])
-            subprocess.check_call(['mount', '-o', 'loop', '--', self.image, self.mountpoint])
+            subprocess.check_call([mkfs, '-q', image])
+            subprocess.check_call(['mount', '-o', 'loop', '--', image, mountpoint])
         except Exception as e:
-            os.remove(self.image)
-            os.rmdir(self.mountpoint)
+            os.rmdir(mountpoint)
+            os.remove(image)
             raise e
 
+        self._mountpoints.append((mountpoint, image))
+        return mountpoint, image
+
+    def setUp(self):
+        self.mountpoint, self.image = self.mount_btrfs()
+
     def tearDown(self):
-        try:
-            subprocess.check_call(['umount', self.mountpoint])
-        finally:
-            os.remove(self.image)
-            os.rmdir(self.mountpoint)
+        for mountpoint, image in self._mountpoints:
+            subprocess.call(['umount', '-R', mountpoint])
+            os.rmdir(mountpoint)
+            os.remove(image)
 
     @staticmethod
     def path_or_fd(path, open_flags=os.O_RDONLY):
@@ -67,4 +106,3 @@ class BtrfsTestCase(unittest.TestCase):
             yield fd
         finally:
             os.close(fd)
-