Merge lp:~salgado/linaro-image-tools/port-create_partitions into lp:linaro-image-tools/11.11

[Guilherme Salgado]

Port most of create_partitions to python

[lp:~salgado/linaro-image-tools/port-create_partitions]

180. By Guilherme Salgado on 2010-12-08

merge trunk

[James Westby (james-w)]

Hi,

Looks fine. I agree sharing more code with do_run would be nice.

The board-specific code in their isn't that nice, but I presume you
are deferring removing that for now?

Thanks,

James

review: Approve

[Guilherme Salgado (salgado)]

I can now use cmd_runner.Popen() here and avoid the duplication. I'll do so after I land that branch.

The board-specific code is going away soon indeed.

Thanks a lot for the review.

[lp:~salgado/linaro-image-tools/port-create_partitions]

181. By Guilherme Salgado on 2010-12-09

merge trunk

182. By Guilherme Salgado on 2010-12-09

Use the new cmd_runner.Popen instead of subprocess.Popen

=== modified file 'hwpack/testing.py'
--- hwpack/testing.py	2010-12-08 00:34:09 +0000
+++ hwpack/testing.py	2010-12-09 20:01:29 +0000
@@ -161,6 +161,15 @@
         fixture.setUp()
         return fixture
 
+    def createTempFileAsFixture(self, prefix='tmp'):
+        """Create a temp file and make sure it is removed on tearDown.
+
+        :return: The filename of the file created.
+        """
+        _, filename = tempfile.mkstemp(prefix=prefix)
+        self.addCleanup(os.unlink, filename)
+        return filename
+
 
 class ConfigFileFixture(object):
 
 
=== modified file 'linaro-media-create'
--- linaro-media-create	2010-12-03 18:55:10 +0000
+++ linaro-media-create	2010-12-09 20:01:29 +0000
@@ -438,40 +438,8 @@
 }
 
 create_partitions() {
-  if [ "${DEVICE}" ]; then
-    sudo parted -s ${DEVICE} mklabel msdos
-  fi
-
-  if [ "${IMAGE_FILE}" ]; then
-    partdev=${IMAGE_FILE}
-  else
-    partdev=${DEVICE}
-  fi
-
-  if [ "$FAT_SIZE" = "32" ]; then
-    PARTITION_TYPE="0x0C"
-  else
-    PARTITION_TYPE="0x0E"
-  fi
-
-  # Create:
-  # - on mx51evk, a one cylinder partition for fixed-offset bootloader data at
-  #   the beginning of the image (size is one cylinder, so 8224768 bytes with
-  #   the first sector for MBR),
-  # - a VFAT or FAT16 partition of 9 cylinders (74027520 bytes, ~70 MiB)
-  #    and a linux partition of the rest
-  # - a Linux type rootfs for the rest of the space
-  (
-    if [ "$DEVIMAGE" = mx51evk ]; then
-      cat <<EOF
-,1,0xDA
-EOF
-    fi
-    cat <<EOF
-,9,$PARTITION_TYPE,*
-,,,-
-EOF
-  ) | sudo sfdisk -D -H $HEADS -S $SECTORS $CYLINDER_ARG $partdev
+  python -m media_create.create_partitions "${DEVIMAGE}" \
+      "${DEVICE-$IMAGE_FILE}" "$FAT_SIZE" "$HEADS" "$SECTORS" "$CYLINDER_ARG"
 
   if [ "${IMAGE_FILE}" ]; then
   VFATOFFSET=$(($(LC_ALL=C fdisk -l -u "$IMAGE_FILE" | grep FAT | awk '{print $3}')*512))
@@ -628,6 +596,7 @@
   fi
 }
 
+# XXX: Apparently there's no need to run qemu-img as root.
 setup_image() {
   sudo qemu-img create -f raw "$IMAGE_FILE" $IMAGE_SIZE
 }
 
=== added file 'media_create/create_partitions.py'
--- media_create/create_partitions.py	1970-01-01 00:00:00 +0000
+++ media_create/create_partitions.py	2010-12-09 20:01:29 +0000
@@ -0,0 +1,87 @@
+import subprocess
+import sys
+
+from media_create import cmd_runner
+
+
+def run_sfdisk_commands(commands, heads, sectors, cylinders_arg, device,
+                        as_root=True):
+    """Run the given commands under sfdisk.
+
+    :param commands: A string of sfdisk commands; each on a separate line.
+    :return: A 2-tuple containing the subprocess' stdout and stderr.
+    """
+    args = ['sfdisk',
+            '-D',
+            '-H', str(heads),
+            '-S', str(sectors)]
+    if cylinders_arg:
+        args.append(cylinders_arg)
+    args.append(device)
+    # XXX: There's some stuff duplicated here from cmd_runner.run() but I
+    # don't see an easy way to consolidate them as a single function, so I'll
+    # leave it for later.
+    if as_root:
+        args = args[:]
+        args.insert(0, 'sudo')
+    proc = cmd_runner.Popen(
+        args, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE)
+    return proc.communicate("%s\n" % commands)
+
+
+def create_partitions(board, device, fat_size, heads, sectors, cylinders_arg):
+    """Partition the given device according to the board requirements.
+
+    :param board: A string with the board type (e.g. beagle, panda, etc)
+    :param device: A string containing the path to the device to partition.
+    :param fat_size: The type of FATs used in the boot partition (16 or 32).
+    :param heads: Number of heads to use in the disk geometry of
+        partitions.
+    :param sectors: Number of sectors to use in the disk geometry of
+        partitions.
+    :param cylinders_arg: A string of the form "-C NN" containing the number
+        of cylinders to use in the disk geometry of partitions.
+    """
+    stdout = []
+    stderr = []
+    is_block_device = device.startswith('/dev/')
+    if is_block_device:
+        # Overwrite any existing partition tables with a fresh one.
+        cmd_runner.run(
+            ['parted', '-s', device, 'mklabel', 'msdos'], as_root=True)
+
+    if fat_size == 32:
+        partition_type = '0x0C'
+    else:
+        partition_type = '0x0E'
+
+    if board == 'mx51evk':
+        # Create a one cylinder partition for fixed-offset bootloader data at
+        # the beginning of the image (size is one cylinder, so 8224768 bytes
+        # with the first sector for MBR).
+        out, err = run_sfdisk_commands(
+            ',1,0xDA', heads, sectors, cylinders_arg, device)
+        stdout.append(out)
+        stderr.append(err)
+
+    # Create a VFAT or FAT16 partition of 9 cylinders (74027520 bytes, ~70
+    # MiB), followed by a Linux-type partition containing the rest of the free
+    # space.
+    sfdisk_cmd = ',9,%s,*\n,,,-' % partition_type
+    out, err = run_sfdisk_commands(
+        sfdisk_cmd, heads, sectors, cylinders_arg, device)
+    stdout.append(out)
+    stderr.append(err)
+    return "\n".join(stdout), "\n".join(stderr)
+
+
+if __name__ == "__main__":
+    board, device, fat_size, heads, sectors, cylinders_arg = sys.argv[1:]
+    fat_size = int(fat_size)
+    heads = int(heads)
+    sectors = int(sectors)
+    stdout, stderr = create_partitions(
+        board, device, fat_size, heads, sectors, cylinders_arg)
+    print stdout
+    print stderr
 
=== modified file 'media_create/tests/fixtures.py'
--- media_create/tests/fixtures.py	2010-12-09 19:17:14 +0000
+++ media_create/tests/fixtures.py	2010-12-09 20:01:29 +0000
@@ -3,6 +3,7 @@
 import subprocess
 import tempfile
 
+from media_create import create_partitions
 from media_create import cmd_runner
 
 
@@ -62,17 +63,11 @@
         setattr(self.obj, self.attr_name, self.orig_attr)
 
 
-class MockDoRun(object):
-    """A mock for do_run() which just stores the args given to it."""
+class MockCmdRunnerPopen(object):
+    """A mock for cmd_runner.Popen() which stores the args given to it."""
     args = None
     def __call__(self, args, **kwargs):
         self.args = args
-        return 0
-
-
-class MockCmdRunnerPopen(object):
-    def __call__(self, args, **kwargs):
-        self.args = args
         self.returncode = 0
         return self
 
@@ -83,7 +78,7 @@
 class MockCmdRunnerPopenFixture(MockSomethingFixture):
     """A test fixture which mocks cmd_runner.do_run with the given mock.
 
-    If no mock is given, MockDoRun is used.
+    If no mock is given, a MockCmdRunnerPopen instance is used.
     """
 
     def __init__(self, mock=None):
@@ -104,3 +99,27 @@
 
     def tearDown(self):
         os.chdir(self.orig_cwd)
+
+
+class MockCallableWithPositionalArgs(object):
+    """A callable mock which just stores the positional args given to it.
+
+    Every time an instance of this is "called", it will append a tuple
+    containing the positional arguments given to it to self.calls.
+    """
+    calls = None
+    return_value = None
+    def __call__(self, *args):
+        if self.calls is None:
+            self.calls = []
+        self.calls.append(args)
+        return self.return_value
+
+
+class MockRunSfdiskCommandsFixture(MockSomethingFixture):
+
+    def __init__(self):
+        mock = MockCallableWithPositionalArgs()
+        mock.return_value = ('', '')
+        super(MockRunSfdiskCommandsFixture, self).__init__(
+            create_partitions, 'run_sfdisk_commands', mock)
 
=== modified file 'media_create/tests/test_media_create.py'
--- media_create/tests/test_media_create.py	2010-12-09 19:17:14 +0000
+++ media_create/tests/test_media_create.py	2010-12-09 20:01:29 +0000
@@ -4,17 +4,19 @@
 import string
 import subprocess
 import sys
-import tempfile
 
 from testtools import TestCase
 
 from hwpack.testing import TestCaseWithFixtures
 
-from media_create.boot_cmd import create_boot_cmd
 from media_create import cmd_runner
 from media_create import ensure_command
-
 from media_create import populate_boot
+from media_create.boot_cmd import create_boot_cmd
+from media_create.create_partitions import (
+    create_partitions,
+    run_sfdisk_commands,
+    )
 from media_create.populate_boot import (
     make_boot_script,
     make_uImage,
@@ -31,6 +33,7 @@
     CreateTarballFixture,
     MockCmdRunnerPopenFixture,
     MockSomethingFixture,
+    MockRunSfdiskCommandsFixture,
     )
 
 
@@ -202,7 +205,7 @@
     def test_get_file_matching(self):
         prefix = ''.join(
             random.choice(string.ascii_lowercase) for x in range(5))
-        file1 = self._create_temp_file_as_fixture(prefix)
+        file1 = self.createTempFileAsFixture(prefix)
         directory = os.path.dirname(file1)
         self.assertEqual(
             file1, _get_file_matching('%s/%s*' % (directory, prefix)))
@@ -210,8 +213,8 @@
     def test_get_file_matching_too_many_files_found(self):
         prefix = ''.join(
             random.choice(string.ascii_lowercase) for x in range(5))
-        file1 = self._create_temp_file_as_fixture(prefix)
-        file2 = self._create_temp_file_as_fixture(prefix)
+        file1 = self.createTempFileAsFixture(prefix)
+        file2 = self.createTempFileAsFixture(prefix)
         directory = os.path.dirname(file1)
         self.assertRaises(
             ValueError, _get_file_matching, '%s/%s*' % (directory, prefix))
@@ -222,23 +225,79 @@
 
     def test_run_mkimage(self):
         # Create a fake boot script.
-        filename = self._create_temp_file_as_fixture()
+        filename = self.createTempFileAsFixture()
         f = open(filename, 'w')
         f.write("setenv bootcmd 'fatload mmc 0:1 0x80000000 uImage;\nboot")
         f.close()
 
-        img = self._create_temp_file_as_fixture()
+        img = self.createTempFileAsFixture()
         # Use that fake boot script to create a boot loader using mkimage.
         # Send stdout to file as mkimage will print to stdout and we don't
         # want that.
         retval = _run_mkimage(
             'script', '0', '0', 'boot script', filename, img,
-            stdout=open(self._create_temp_file_as_fixture(), 'w'),
+            stdout=open(self.createTempFileAsFixture(), 'w'),
             as_root=False)
 
         self.assertEqual(0, retval)
 
-    def _create_temp_file_as_fixture(self, prefix='tmp'):
-        _, filename = tempfile.mkstemp(prefix=prefix)
-        self.addCleanup(os.unlink, filename)
-        return filename
+
+class TestCreatePartitions(TestCaseWithFixtures):
+
+    def test_create_partitions_for_mx51evk(self):
+        # For this board we create a one cylinder partition at the beginning.
+        popen_fixture = self.useFixture(MockCmdRunnerPopenFixture())
+        sfdisk_fixture = self.useFixture(MockRunSfdiskCommandsFixture())
+
+        create_partitions('mx51evk', '/dev/sdz', 32, 255, 63, '')
+
+        self.assertEqual(
+            ['sudo', 'parted', '-s', '/dev/sdz', 'mklabel', 'msdos'],
+            popen_fixture.mock.args)
+        self.assertEqual(
+            [(',1,0xDA', 255, 63, '', '/dev/sdz'),
+             (',9,0x0C,*\n,,,-', 255, 63, '', '/dev/sdz')],
+            sfdisk_fixture.mock.calls)
+
+    def test_create_partitions_for_beagle(self):
+        popen_fixture = self.useFixture(MockCmdRunnerPopenFixture())
+        sfdisk_fixture = self.useFixture(MockRunSfdiskCommandsFixture())
+
+        create_partitions('beagle', '/dev/sdz', 32, 255, 63, '')
+
+        self.assertEqual(
+            ['sudo', 'parted', '-s', '/dev/sdz', 'mklabel', 'msdos'],
+            popen_fixture.mock.args)
+        self.assertEqual(
+            [(',9,0x0C,*\n,,,-', 255, 63, '', '/dev/sdz')],
+            sfdisk_fixture.mock.calls)
+
+    def test_create_partitions_with_img_file(self):
+        popen_fixture = self.useFixture(MockCmdRunnerPopenFixture())
+        sfdisk_fixture = self.useFixture(MockRunSfdiskCommandsFixture())
+
+        tempfile = self.createTempFileAsFixture()
+        create_partitions('beagle', tempfile, 32, 255, 63, '')
+
+        # popen() was not called as there's no existing partition table for
+        # us to overwrite on the image file.
+        self.assertEqual(None, popen_fixture.mock.args)
+
+        self.assertEqual(
+            [(',9,0x0C,*\n,,,-', 255, 63, '', tempfile)],
+            sfdisk_fixture.mock.calls)
+
+    def test_run_sfdisk_commands(self):
+        tempfile = self.createTempFileAsFixture()
+        cmd_runner.run(['qemu-img', 'create', '-f', 'raw', tempfile, '10M'],
+                       stdout=subprocess.PIPE)
+        stdout, stderr = run_sfdisk_commands(
+            ',1,0xDA', 5, 63, '', tempfile, as_root=False)
+        self.assertIn('Successfully wrote the new partition table', stdout)
+
+    def test_run_sfdisk_commands_raises_on_non_zero_returncode(self):
+        tempfile = self.createTempFileAsFixture()
+        self.assertRaises(
+            cmd_runner.SubcommandNonZeroReturnValue,
+            run_sfdisk_commands,
+            ',1,0xDA', 5, 63, '', tempfile, as_root=False)