=== modified file 'linaro_image_tools/media_create/rootfs.py'
--- linaro_image_tools/media_create/rootfs.py	2011-07-18 14:38:39 +0000
+++ linaro_image_tools/media_create/rootfs.py	2011-07-21 21:31:15 +0000
@@ -17,19 +17,27 @@
 # You should have received a copy of the GNU General Public License
 # along with Linaro Image Tools.  If not, see <http://www.gnu.org/licenses/>.
 
-import glob
+import atexit
 import os
 import subprocess
 import tempfile
 
 from linaro_image_tools import cmd_runner
 
+
+def umount(device):
+    # The old code used to ignore failures here, but I don't think that's
+    # desirable so I'm using cmd_runner.run()'s standard behaviour, which will
+    # fail on a non-zero return value.
+    cmd_runner.run(['umount', device], as_root=True).wait()
+
+
 def populate_partition(content_dir, root_disk, partition):
     os.makedirs(root_disk)
     cmd_runner.run(['mount', partition, root_disk], as_root=True).wait()
+    atexit.register(umount, root_disk)
     move_contents(content_dir, root_disk)
     cmd_runner.run(['sync']).wait()
-    cmd_runner.run(['umount', root_disk], as_root=True).wait()
 
 
 def rootfs_mount_options(rootfs_type):
@@ -49,11 +57,11 @@
     This consists of:
       1. Create a directory on the path specified by root_disk
       2. Mount the given partition onto the created directory.
-      3. Move the contents of content_dir to that directory.
-      4. If should_create_swap, then create it with the given size.
-      5. Add fstab entries for the / filesystem and swap (if created).
-      6. Create a /etc/flash-kernel.conf containing the target's boot device.
-      7. Unmount the partition we mounted on step 2.
+      3. Setup an atexit handler to unmount the partition mounted above.
+      4. Move the contents of content_dir to that directory.
+      5. If should_create_swap, then create it with the given size.
+      6. Add fstab entries for the / filesystem and swap (if created).
+      7. Create a /etc/flash-kernel.conf containing the target's boot device.
     """
     print "\nPopulating rootfs partition"
     print "Be patient, this may take a few minutes\n"
@@ -61,6 +69,9 @@
     os.makedirs(root_disk)
 
     cmd_runner.run(['mount', partition, root_disk], as_root=True).wait()
+    # We use an atexit handler to umount the partition to make sure it's
+    # umounted even if something fails when it's being populated.
+    atexit.register(umount, root_disk)
 
     move_contents(content_dir, root_disk)
 
@@ -91,10 +102,6 @@
     create_flash_kernel_config(root_disk, 1 + partition_offset)
 
     cmd_runner.run(['sync']).wait()
-    # The old code used to ignore failures here, but I don't think that's
-    # desirable so I'm using cmd_runner.run()'s standard behaviour, which will
-    # fail on a non-zero return value.
-    cmd_runner.run(['umount', root_disk], as_root=True).wait()
 
 
 def create_flash_kernel_config(root_disk, boot_partition_number):

=== modified file 'linaro_image_tools/media_create/tests/test_media_create.py'
--- linaro_image_tools/media_create/tests/test_media_create.py	2011-07-19 22:19:43 +0000
+++ linaro_image_tools/media_create/tests/test_media_create.py	2011-07-21 21:24:13 +0000
@@ -101,6 +101,7 @@
     move_contents,
     populate_rootfs,
     rootfs_mount_options,
+    umount,
     write_data_to_protected_file,
     )
 from linaro_image_tools.media_create.tests.fixtures import (
@@ -2104,6 +2105,8 @@
         def fake_create_flash_kernel_config(disk, partition_offset):
             self.create_flash_kernel_config_called = True
 
+        atexit_fixture = self.useFixture(MockSomethingFixture(
+            atexit, 'register', AtExitRegister()))
         # Mock stdout, cmd_runner.Popen(), append_to_fstab and
         # create_flash_kernel_config.
         self.useFixture(MockSomethingFixture(
@@ -2149,9 +2152,11 @@
             '%s dd if=/dev/zero of=%s bs=1M count=100' % (
                sudo_args, swap_file),
             '%s mkswap %s' % (sudo_args, swap_file),
-            'sync',
-            '%s umount %s' % (sudo_args, root_disk)]
+            'sync']
         self.assertEqual(expected, popen_fixture.mock.commands_executed)
+        self.assertEqual(1, len(atexit_fixture.mock.funcs))
+        self.assertEqual(
+            (umount, (root_disk,), {}), atexit_fixture.mock.funcs[0])
 
     def test_create_flash_kernel_config(self):
         fixture = self.useFixture(MockCmdRunnerPopenFixture())

