Refactor remount() into common function

common/fs/__init__.py

@@ -14,3 +14,25 @@ def get_partitions():
 				raise RuntimeError('Unable to parse {line} in {path}'.format(line=line, path=path))
 			matches[match.group('dev_name')] = match.groupdict()
 	return matches
+
+
+def remount(volume, fn):
+	from base.fs.partitionmaps.none import NoPartitions
+
+	p_map = volume.partition_map
+	volume.unmount_specials()
+	if hasattr(p_map, 'boot'):
+		boot_dir = p_map.boot.mount_dir
+		p_map.boot.unmount()
+	root_dir = p_map.root.mount_dir
+	p_map.root.unmount()
+	if not isinstance(p_map, NoPartitions):
+		p_map.unmap(volume)
+		fn()
+		p_map.map(volume)
+	else:
+		fn()
+	p_map.root.mount(root_dir)
+	if hasattr(p_map, 'boot'):
+		p_map.boot.mount(boot_dir)
+	volume.mount_specials()

plugins/prebootstrapped/tasks.py

@@ -48,7 +48,10 @@ class CopyImage(Task):
 	def run(self, info):
 		loopback_backup_name = 'volume-{id:x}.{ext}.backup'.format(id=info.run_id, ext=info.volume.extension)
 		destination = os.path.join(info.manifest.bootstrapper['workspace'], loopback_backup_name)
-		copyfile(info.volume.image_path, destination)
+
+		def mk_snapshot():
+			copyfile(info.volume.image_path, destination)
+		remount(info.volume, mk_snapshot)
 		msg = 'A copy of the bootstrapped volume was created. Path: {path}'.format(path=destination)
 		log.info(msg)
 

providers/virtualbox/tasks/boot.py

@@ -18,31 +18,23 @@ class ConfigureGrub(Task):
 
 		from base.fs.partitionmaps.none import NoPartitions
 		from base.fs.partitionmaps.gpt import GPTPartitionMap
+		from common.fs import remount
 		p_map = info.volume.partition_map
 
-		def remount(fn):
+		def mk_remount_fn(fn):
-			# GRUB cannot deal with installing to loopback devices
+			def set_device_path():
-			# so we fake a real harddisk with dmsetup.
-			# Guide here: http://ebroder.net/2009/08/04/installing-grub-onto-a-disk-image/
-			info.volume.unmount_specials()
-			if hasattr(p_map, 'boot'):
-				boot_dir = p_map.boot.mount_dir
-				p_map.boot.unmount()
-			p_map.root.unmount()
-			if not isinstance(p_map, NoPartitions):
-				p_map.unmap(info.volume)
 				fn()
-				p_map.map(info.volume)
+				if isinstance(p_map, NoPartitions):
-			else:
+					p_map.root.device_path = info.volume.device_path
-				fn()
+			return set_device_path
-				p_map.root.device_path = info.volume.device_path
+		link_fn = mk_remount_fn(info.volume.link_dm_node)
-			p_map.root.mount(info.root)
+		unlink_fn = mk_remount_fn(info.volume.unlink_dm_node)
-			if hasattr(p_map, 'boot'):
-				p_map.boot.mount(boot_dir)
-			info.volume.mount_specials()
 
+		# GRUB cannot deal with installing to loopback devices
+		# so we fake a real harddisk with dmsetup.
+		# Guide here: http://ebroder.net/2009/08/04/installing-grub-onto-a-disk-image/
 		if isinstance(info.volume, LoopbackVolume):
-			remount(info.volume.link_dm_node)
+			remount(info.volume, link_fn)
 		try:
 			[device_path] = log_check_call(['readlink', '-f', info.volume.device_path])
 			device_map_path = os.path.join(grub_dir, 'device.map')
@@ -66,8 +58,8 @@ class ConfigureGrub(Task):
 			log_check_call(['/usr/sbin/chroot', info.root, '/usr/sbin/update-grub'])
 		except Exception as e:
 			if isinstance(info.volume, LoopbackVolume):
-				remount(info.volume.unlink_dm_node)
+				remount(info.volume, unlink_fn)
 			raise e
 
 		if isinstance(info.volume, LoopbackVolume):
-			remount(info.volume.unlink_dm_node)
+			remount(info.volume, unlink_fn)