1 files changed, 30 insertions, 35 deletions
diff --git a/morphlib/exts/rawdisk.write b/morphlib/exts/rawdisk.write
index e1a75fe0..b17f8aa7 100755
--- a/morphlib/exts/rawdisk.write
+++ b/morphlib/exts/rawdisk.write
@@ -67,51 +67,46 @@ class RawDiskWriteExtension(morphlib.writeexts.WriteExtension):
     def upgrade_local_system(self, raw_disk, temp_root):
         self.complete_fstab_for_btrfs_layout(temp_root)
 
-        mp = self.mount(raw_disk)
+        with self.mount(raw_disk) as mp:
+            version_label = self.get_version_label(mp)
+            self.status(msg='Updating image to a new version with label %s' %
+                        version_label)
+
+            version_root = os.path.join(mp, 'systems', version_label)
+            os.mkdir(version_root)
+
+            old_orig = os.path.join(mp, 'systems', 'default', 'orig')
+            new_orig = os.path.join(version_root, 'orig')
+            cliapp.runcmd(
+                ['btrfs', 'subvolume', 'snapshot', old_orig, new_orig])
+
+            cliapp.runcmd(
+                ['rsync', '-a', '--checksum', '--numeric-ids', '--delete',
+                 temp_root + os.path.sep, new_orig])
+
+            self.create_run(version_root)
+
+            default_path = os.path.join(mp, 'systems', 'default')
+            if os.path.exists(default_path):
+                os.remove(default_path)
+            else:
+                # we are upgrading and old system that does
+                # not have an updated extlinux config file
+                if self.bootloader_config_is_wanted():
+                    self.generate_bootloader_config(mp)
+                self.install_bootloader(mp)
+            os.symlink(version_label, default_path)
 
-        version_label = self.get_version_label(mp)
-        self.status(msg='Updating image to a new version with label %s' %
-                    version_label)
-
-        version_root = os.path.join(mp, 'systems', version_label)
-        os.mkdir(version_root)
-
-        old_orig = os.path.join(mp, 'systems', 'default', 'orig')
-        new_orig = os.path.join(version_root, 'orig')
-        cliapp.runcmd(
-            ['btrfs', 'subvolume', 'snapshot', old_orig, new_orig])
-
-        cliapp.runcmd(
-            ['rsync', '-a', '--checksum', '--numeric-ids', '--delete',
-             temp_root + os.path.sep, new_orig])
-
-        self.create_run(version_root)
-
-        default_path = os.path.join(mp, 'systems', 'default')
-        if os.path.exists(default_path):
-            os.remove(default_path)
-        else:
-            # we are upgrading and old system that does
-            # not have an updated extlinux config file
             if self.bootloader_config_is_wanted():
-                self.generate_bootloader_config(mp)
-            self.install_bootloader(mp)
-        os.symlink(version_label, default_path)
-
-        if self.bootloader_config_is_wanted():
-            self.install_kernel(version_root, temp_root)
-
-        self.unmount(mp)
+                self.install_kernel(version_root, temp_root)
 
     def get_version_label(self, mp):
         version_label = os.environ.get('VERSION_LABEL')
 
         if version_label is None:
-            self.unmount(mp)
             raise cliapp.AppException('VERSION_LABEL was not given')
 
         if os.path.exists(os.path.join(mp, 'systems', version_label)):
-            self.unmount(mp)
             raise cliapp.AppException('VERSION_LABEL %s already exists'
                                       % version_label)