update_payload: De-duplicate paycheck partition handling for major version 2

This commit makes handling of major version 2 rootfs/kernel partitions more
generic, with the goal of reusing the same code paths for major version 2
(which has an arbitrary number of partitions per payload).

BUG=b:794404
TEST=no errors during run_unittests and test_paycheck.sh

Change-Id: Ic6e8cdaae557c10dfbd302dba0e52ff8b2870c9f
Reviewed-on: https://chromium-review.googlesource.com/1101500
Commit-Ready: Tudor Brindus <tbrindus@chromium.org>
Tested-by: Tudor Brindus <tbrindus@chromium.org>
Reviewed-by: Amin Hassani <ahassani@chromium.org>
diff --git a/scripts/update_payload/checker_unittest.py b/scripts/update_payload/checker_unittest.py
index 68f1807..d42f4b4 100755
--- a/scripts/update_payload/checker_unittest.py
+++ b/scripts/update_payload/checker_unittest.py
@@ -485,13 +485,16 @@
                    fail_bad_nki or fail_bad_nri or fail_old_kernel_fs_size or
                    fail_old_rootfs_fs_size or fail_new_kernel_fs_size or
                    fail_new_rootfs_fs_size)
+    part_sizes = {
+        common.ROOTFS: rootfs_part_size,
+        common.KERNEL: kernel_part_size
+    }
+
     if should_fail:
       self.assertRaises(PayloadError, payload_checker._CheckManifest, report,
-                        rootfs_part_size, kernel_part_size)
+                        part_sizes)
     else:
-      self.assertIsNone(payload_checker._CheckManifest(report,
-                                                       rootfs_part_size,
-                                                       kernel_part_size))
+      self.assertIsNone(payload_checker._CheckManifest(report, part_sizes))
 
   def testCheckLength(self):
     """Tests _CheckLength()."""
@@ -1085,7 +1088,10 @@
     report = checker._PayloadReport()
 
     # We have to check the manifest first in order to set signature attributes.
-    payload_checker._CheckManifest(report, rootfs_part_size, kernel_part_size)
+    payload_checker._CheckManifest(report, {
+        common.ROOTFS: rootfs_part_size,
+        common.KERNEL: kernel_part_size
+    })
 
     should_fail = (fail_empty_sigs_blob or fail_missing_pseudo_op or
                    fail_mismatched_pseudo_op or fail_sig_missing_fields or