=== modified file 'linaro_image_tools/media_create/android_boards.py'
--- linaro_image_tools/media_create/android_boards.py	2011-07-18 18:25:30 +0000
+++ linaro_image_tools/media_create/android_boards.py	2011-07-25 16:48:37 +0000
@@ -78,6 +78,8 @@
     @classmethod
     def populate_boot_script(cls, boot_partition, boot_disk, consoles):
         cmd_runner.run(['mkdir', '-p', boot_disk]).wait()
+        # TODO: Use partition_mounted() here to make sure the partition is
+        # always umounted after we're done.
         cmd_runner.run(['mount', boot_partition, boot_disk],
             as_root=True).wait()
 

=== modified file 'linaro_image_tools/media_create/boards.py'
--- linaro_image_tools/media_create/boards.py	2011-07-19 22:19:43 +0000
+++ linaro_image_tools/media_create/boards.py	2011-07-25 16:48:37 +0000
@@ -37,7 +37,8 @@
 
 from linaro_image_tools import cmd_runner
 
-from linaro_image_tools.media_create.partitions import SECTOR_SIZE
+from linaro_image_tools.media_create.partitions import (
+    partition_mounted, SECTOR_SIZE)
 
 
 KERNEL_GLOB = 'vmlinuz-*-%(kernel_flavor)s'
@@ -465,28 +466,22 @@
         uboot_parts_dir = os.path.join(chroot_dir, parts_dir)
 
         cmd_runner.run(['mkdir', '-p', boot_disk]).wait()
-        cmd_runner.run(['mount', boot_partition, boot_disk],
-            as_root=True).wait()
-
-        if cls.uboot_in_boot_part:
-            assert cls.uboot_flavor is not None, (
-                "uboot_in_boot_part is set but not uboot_flavor")
-            with cls.hardwarepack_handler:
-                uboot_bin = cls.get_file('u_boot', default=os.path.join(
+        with partition_mounted(boot_partition, boot_disk):
+            if cls.uboot_in_boot_part:
+                assert cls.uboot_flavor is not None, (
+                    "uboot_in_boot_part is set but not uboot_flavor")
+                with cls.hardwarepack_handler:
+                    default = os.path.join(
                         chroot_dir, 'usr', 'lib', 'u-boot', cls.uboot_flavor,
-                        'u-boot.bin'))
-                cmd_runner.run(
-                    ['cp', '-v', uboot_bin, boot_disk], as_root=True).wait()
-
-        cls.make_boot_files(
-            uboot_parts_dir, is_live, is_lowmem, consoles, chroot_dir,
-            rootfs_uuid, boot_disk, boot_device_or_file)
-
-        cmd_runner.run(['sync']).wait()
-        try:
-            cmd_runner.run(['umount', boot_disk], as_root=True).wait()
-        except cmd_runner.SubcommandNonZeroReturnValue:
-            pass
+                        'u-boot.bin')
+                    uboot_bin = cls.get_file('u_boot', default=default)
+                    proc = cmd_runner.run(
+                        ['cp', '-v', uboot_bin, boot_disk], as_root=True)
+                    proc.wait()
+
+            cls.make_boot_files(
+                uboot_parts_dir, is_live, is_lowmem, consoles, chroot_dir,
+                rootfs_uuid, boot_disk, boot_device_or_file)
 
     @classmethod
     def _get_kflavor_files(cls, path):

=== modified file 'linaro_image_tools/media_create/chroot_utils.py'
--- linaro_image_tools/media_create/chroot_utils.py	2011-04-29 11:47:02 +0000
+++ linaro_image_tools/media_create/chroot_utils.py	2011-07-25 16:48:37 +0000
@@ -93,6 +93,8 @@
     prepare_chroot(chroot_dir, tmp_dir)
 
     try:
+        # TODO: Use the partition_mounted() contextmanager here and get rid of
+        # mount_chroot_proc() altogether.
         mount_chroot_proc(chroot_dir)
         print "-" * 60
         print "Installing (apt-get) %s in target rootfs." % " ".join(packages)

=== modified file 'linaro_image_tools/media_create/partitions.py'
--- linaro_image_tools/media_create/partitions.py	2011-06-29 09:16:39 +0000
+++ linaro_image_tools/media_create/partitions.py	2011-07-26 12:10:08 +0000
@@ -18,7 +18,9 @@
 # along with Linaro Image Tools.  If not, see <http://www.gnu.org/licenses/>.
 
 import atexit
+from contextlib import contextmanager
 import glob
+import logging
 import re
 import subprocess
 import time
@@ -169,6 +171,39 @@
     return bootfs, rootfs
 
 
+def umount(path):
+    # 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', path], as_root=True).wait()
+
+
+@contextmanager
+def partition_mounted(device, path, *args):
+    """A context manager that mounts the given device and umounts when done.
+
+    We use a try/finally to make sure the device is umounted even if there's
+    an uncaught exception in the with block.  Also, before umounting we call
+    'sync'.
+
+    :param *args: Extra arguments to the mount command.
+    """
+    subprocess_args = ['mount', device, path]
+    subprocess_args.extend(args)
+    cmd_runner.run(subprocess_args, as_root=True).wait()
+    try:
+        yield
+    finally:
+        cmd_runner.run(['sync']).wait()
+        try:
+            umount(path)
+        except cmd_runner.SubcommandNonZeroReturnValue, e:
+            logger = logging.getLogger("linaro_image_tools")
+            logger.warn("Failed to umount %s, but ignoring it because of a "
+                        "previous error" % path)
+            logger.warn(e)
+
+
 def get_uuid(partition):
     """Find UUID of the given partition."""
     proc = cmd_runner.run(

=== modified file 'linaro_image_tools/media_create/rootfs.py'
--- linaro_image_tools/media_create/rootfs.py	2011-07-21 21:31:15 +0000
+++ linaro_image_tools/media_create/rootfs.py	2011-07-25 16:48:37 +0000
@@ -17,27 +17,19 @@
 # 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 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()
+from linaro_image_tools.media_create.partitions import partition_mounted
 
 
 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()
+    with partition_mounted(partition, root_disk):
+        move_contents(content_dir, root_disk)
 
 
 def rootfs_mount_options(rootfs_type):
@@ -68,40 +60,34 @@
     # Create a directory to mount the rootfs partition.
     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)
-
-    mount_options = rootfs_mount_options(rootfs_type)
-    fstab_additions = ["UUID=%s / %s  %s 0 1" % (
-            rootfs_uuid, rootfs_type, mount_options)]
-    if should_create_swap:
-        print "\nCreating SWAP File\n"
-        if has_space_left_for_swap(root_disk, swap_size):
-            proc = cmd_runner.run([
-                'dd',
-                'if=/dev/zero',
-                'of=%s/SWAP.swap' % root_disk,
-                'bs=1M',
-                'count=%s' % swap_size], as_root=True)
-            proc.wait()
-            proc = cmd_runner.run(
-                ['mkswap', '%s/SWAP.swap' % root_disk], as_root=True)
-            proc.wait()
-            fstab_additions.append("/SWAP.swap  none  swap  sw  0 0")
-        else:
-            print ("Swap file is bigger than space left on partition; "
-                   "continuing without swap.")
-
-    append_to_fstab(root_disk, fstab_additions)
-
-    print "\nCreating /etc/flash-kernel.conf\n"
-    create_flash_kernel_config(root_disk, 1 + partition_offset)
-
-    cmd_runner.run(['sync']).wait()
+    with partition_mounted(partition, root_disk):
+        move_contents(content_dir, root_disk)
+
+        mount_options = rootfs_mount_options(rootfs_type)
+        fstab_additions = ["UUID=%s / %s  %s 0 1" % (
+                rootfs_uuid, rootfs_type, mount_options)]
+        if should_create_swap:
+            print "\nCreating SWAP File\n"
+            if has_space_left_for_swap(root_disk, swap_size):
+                proc = cmd_runner.run([
+                    'dd',
+                    'if=/dev/zero',
+                    'of=%s/SWAP.swap' % root_disk,
+                    'bs=1M',
+                    'count=%s' % swap_size], as_root=True)
+                proc.wait()
+                proc = cmd_runner.run(
+                    ['mkswap', '%s/SWAP.swap' % root_disk], as_root=True)
+                proc.wait()
+                fstab_additions.append("/SWAP.swap  none  swap  sw  0 0")
+            else:
+                print ("Swap file is bigger than space left on partition; "
+                       "continuing without swap.")
+
+        append_to_fstab(root_disk, fstab_additions)
+
+        print "\nCreating /etc/flash-kernel.conf\n"
+        create_flash_kernel_config(root_disk, 1 + partition_offset)
 
 
 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-21 21:24:13 +0000
+++ linaro_image_tools/media_create/tests/test_media_create.py	2011-07-26 12:10:08 +0000
@@ -89,6 +89,7 @@
     get_android_loopback_devices,
     get_boot_and_root_partitions_for_media,
     Media,
+    partition_mounted,
     run_sfdisk_commands,
     setup_partitions,
     get_uuid,
@@ -101,7 +102,6 @@
     move_contents,
     populate_rootfs,
     rootfs_mount_options,
-    umount,
     write_data_to_protected_file,
     )
 from linaro_image_tools.media_create.tests.fixtures import (
@@ -2016,6 +2016,52 @@
             popen_fixture.mock.commands_executed)
 
 
+class TestException(Exception):
+    """Just a test exception."""
+
+
+class TestMountedPartitionContextManager(TestCaseWithFixtures):
+
+    def test_basic(self):
+        popen_fixture = self.useFixture(MockCmdRunnerPopenFixture())
+        def test_func():
+            with partition_mounted('foo', 'bar', '-t', 'proc'):
+                pass
+        test_func()
+        expected = ['%s mount foo bar -t proc' % sudo_args,
+                    'sync',
+                    '%s umount bar' % sudo_args]
+        self.assertEqual(expected, popen_fixture.mock.commands_executed)
+
+    def test_exception_raised_inside_with_block(self):
+        popen_fixture = self.useFixture(MockCmdRunnerPopenFixture())
+        def test_func():
+            with partition_mounted('foo', 'bar'):
+                raise TestException('something')
+        try:
+            test_func()
+        except TestException:
+            pass
+        expected = ['%s mount foo bar' % sudo_args,
+                    'sync',
+                    '%s umount bar' % sudo_args]
+        self.assertEqual(expected, popen_fixture.mock.commands_executed)
+
+    def test_umount_failure(self):
+        # We ignore a SubcommandNonZeroReturnValue from umount because
+        # otherwise it could shadow an exception raised inside the 'with'
+        # block.
+        popen_fixture = self.useFixture(MockCmdRunnerPopenFixture())
+        def failing_umount(path):
+            raise cmd_runner.SubcommandNonZeroReturnValue('umount', 1)
+        self.useFixture(MockSomethingFixture(
+            partitions, 'umount', failing_umount))
+        def test_func():
+            with partition_mounted('foo', 'bar'):
+                pass
+        test_func()
+
+
 class TestPopulateBoot(TestCaseWithFixtures):
 
     expected_args = (
@@ -2105,8 +2151,6 @@
         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(
@@ -2152,11 +2196,9 @@
             '%s dd if=/dev/zero of=%s bs=1M count=100' % (
                sudo_args, swap_file),
             '%s mkswap %s' % (sudo_args, swap_file),
-            'sync']
+            'sync',
+            '%s umount %s' % (sudo_args, root_disk)]
         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())

