From 5aa74fd32200d84e81e1473eed9e155ebd5a379b Mon Sep 17 00:00:00 2001
From: Pedro Alvarez <pedro.alvarez@codethink.co.uk>
Date: Tue, 11 Nov 2014 15:42:07 +0000
Subject: Split create_local_system on various methods

This way we can still use create_local_system to create
a raw disk, but also reuse bits of it to be able e.g. to
deploy to devices.
---
 writeexts.py | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/writeexts.py b/writeexts.py
index 31e7a0b6..7030c0b5 100644
--- a/writeexts.py
+++ b/writeexts.py
@@ -24,6 +24,7 @@ import time
 import tempfile
 import errno
 import stat
+import contextlib
 
 import morphlib
 
@@ -147,25 +148,39 @@ class WriteExtension(cliapp.Application):
 
     def create_local_system(self, temp_root, raw_disk):
         '''Create a raw system image locally.'''
+
+        with self.created_disk_image(raw_disk):
+            self.format_btrfs(raw_disk)
+            self.create_system(temp_root, raw_disk)
+
+    @contextlib.contextmanager
+    def created_disk_image(self, location):
         size = self.get_disk_size()
         if not size:
             raise cliapp.AppException('DISK_SIZE is not defined')
-        self.create_raw_disk_image(raw_disk, size)
+        self.create_raw_disk_image(location, size)
+        try:
+            yield
+        except BaseException:
+            os.unlink(location)
+            raise
+
+    def format_btrfs(self, raw_disk):
         try:
             self.mkfs_btrfs(raw_disk)
-            mp = self.mount(raw_disk)
         except BaseException:
             sys.stderr.write('Error creating disk image')
-            os.remove(raw_disk)
             raise
+
+    def create_system(self, temp_root, raw_disk):
         try:
+            mp = self.mount(raw_disk)
             self.create_btrfs_system_layout(
                 temp_root, mp, version_label='factory',
                 disk_uuid=self.get_uuid(raw_disk))
         except BaseException, e:
             sys.stderr.write('Error creating Btrfs system layout')
             self.unmount(mp)
-            os.remove(raw_disk)
             raise
         else:
             self.unmount(mp)
-- 
cgit v1.2.1