libbtrfsutil: allow tests to create multiple Btrfs instances

Some upcoming tests will need to create a second Btrfs filesystem, so
add support for this to the test helpers.

Signed-off-by: Omar Sandoval <osandov@fb.com>
Signed-off-by: David Sterba <dsterba@suse.com>

1 files changed, 22 insertions, 13 deletions

diff --git a/libbtrfsutil/python/tests/__init__.py b/libbtrfsutil/python/tests/__init__.py
index 4bc11990..9fd6f6de 100644
--- a/libbtrfsutil/python/tests/__init__.py
+++ b/libbtrfsutil/python/tests/__init__.py
@@ -57,14 +57,18 @@ def regain_privs():
 
 @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'):
@@ -72,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):