cros payload: support show operations in version 2.

Use operations in partitions instead of install_operations and
kernel_install_operations if payload major version is 2.

BUG=None
TEST=./cros_payload_unittest
CQ-DEPEND=CL:299498

Change-Id: I764454ddaa12e55096275d289a328f07cb7dd55a
Reviewed-on: https://chromium-review.googlesource.com/299743
Commit-Ready: Sen Jiang <senj@chromium.org>
Tested-by: Sen Jiang <senj@chromium.org>
Reviewed-by: Alex Deymo <deymo@chromium.org>

2 files changed, 117 insertions, 25 deletions

diff --git a/cli/cros/cros_payload.py b/cli/cros/cros_payload.py
index ebac811ab..61c41417a 100644
--- a/cli/cros/cros_payload.py
+++ b/cli/cros/cros_payload.py
@@ -18,6 +18,8 @@ from chromite.cli import command
 # Needed for the dev.host.lib import below.
 sys.path.insert(0, os.path.join(constants.SOURCE_ROOT, 'src', 'platform'))
 
+MAJOR_PAYLOAD_VERSION_CHROMEOS = 1
+MAJOR_PAYLOAD_VERSION_BRILLO = 2
 
 def DisplayValue(key, value):
   """Print out a key, value pair with values left-aligned."""
@@ -85,9 +87,15 @@ Example:
   def _DisplayManifest(self):
     """Show information from the payload manifest."""
     manifest = self.payload.manifest
-    DisplayValue('Number of operations', len(manifest.install_operations))
-    DisplayValue('Number of kernel ops',
-                 len(manifest.kernel_install_operations))
+    if self.payload.header.version == MAJOR_PAYLOAD_VERSION_BRILLO:
+      DisplayValue('Number of partitions', len(manifest.partitions))
+      for partition in manifest.partitions:
+        DisplayValue('  Number of "%s" ops' % partition.partition_name,
+                     len(partition.operations))
+    else:
+      DisplayValue('Number of operations', len(manifest.install_operations))
+      DisplayValue('Number of kernel ops',
+                   len(manifest.kernel_install_operations))
     DisplayValue('Block size', manifest.block_size)
     DisplayValue('Minor version', manifest.minor_version)
 
@@ -162,8 +170,12 @@ Example:
     read_blocks = 0
     written_blocks = 0
     num_write_seeks = 0
-    for operations in (manifest.install_operations,
-                       manifest.kernel_install_operations):
+    if self.payload.header.version == MAJOR_PAYLOAD_VERSION_BRILLO:
+      partitions_operations = [part.operations for part in manifest.partitions]
+    else:
+      partitions_operations = [manifest.install_operations,
+                               manifest.kernel_install_operations]
+    for operations in partitions_operations:
       last_ext = None
       for curr_op in operations:
         read_blocks += sum([ext.num_blocks for ext in curr_op.src_extents])
@@ -207,7 +219,12 @@ Example:
       self._DisplayStats(self.payload.manifest)
     if self.options.list_ops:
       print()
-      self._DisplayOps('Install operations',
-                       self.payload.manifest.install_operations)
-      self._DisplayOps('Kernel install operations',
-                       self.payload.manifest.kernel_install_operations)
+      if self.payload.header.version == MAJOR_PAYLOAD_VERSION_BRILLO:
+        for partition in self.payload.manifest.partitions:
+          self._DisplayOps('%s install operations' % partition.partition_name,
+                           partition.operations)
+      else:
+        self._DisplayOps('Install operations',
+                         self.payload.manifest.install_operations)
+        self._DisplayOps('Kernel install operations',
+                         self.payload.manifest.kernel_install_operations)
diff --git a/cli/cros/cros_payload_unittest.py b/cli/cros/cros_payload_unittest.py
index d99f0e9a1..f8fd78d20 100644
--- a/cli/cros/cros_payload_unittest.py
+++ b/cli/cros/cros_payload_unittest.py
@@ -50,10 +50,17 @@ class FakeOp(object):
   def HasField(self, field):
     return hasattr(self, field)
 
+class FakePartition(object):
+  """Fake PartitionUpdate field for testing."""
+
+  def __init__(self, partition_name, operations):
+    self.partition_name = partition_name
+    self.operations = operations
+
 class FakeManifest(object):
   """Fake manifest for testing."""
 
-  def __init__(self):
+  def __init__(self, major_version):
     FakeExtent = collections.namedtuple('FakeExtent',
                                         ['start_block', 'num_blocks'])
     self.install_operations = [FakeOp([],
@@ -67,6 +74,11 @@ class FakeManifest(object):
         [FakeExtent(x, x) for x in xrange(20)],
         update_payload.common.OpType.SOURCE_COPY,
         src_length=4096)]
+    if major_version == cros_payload.MAJOR_PAYLOAD_VERSION_BRILLO:
+      self.partitions = [FakePartition('rootfs', self.install_operations),
+                         FakePartition('kernel',
+                                       self.kernel_install_operations)]
+      self.install_operations = self.kernel_install_operations = []
     self.block_size = 4096
     self.minor_version = 4
     FakePartInfo = collections.namedtuple('FakePartInfo', ['size'])
@@ -84,12 +96,12 @@ class FakeManifest(object):
 class FakePayload(object):
   """Fake payload for testing."""
 
-  def __init__(self):
+  def __init__(self, major_version):
     FakeHeader = collections.namedtuple('FakeHeader',
                                         ['version', 'manifest_len'])
-    self._header = FakeHeader('111', 222)
+    self._header = FakeHeader(major_version, 222)
     self.header = None
-    self._manifest = FakeManifest()
+    self._manifest = FakeManifest(major_version)
     self.manifest = None
 
     self._blobs = {}
@@ -137,13 +149,14 @@ class PayloadCommandTest(cros_test_lib.MockOutputTestCase):
   def testRun(self):
     """Verify that Run parses and displays the payload like we expect."""
     payload_cmd = cros_payload.PayloadCommand(FakeOption(action='show'))
-    self.PatchObject(update_payload, 'Payload', return_value=FakePayload())
+    self.PatchObject(update_payload, 'Payload', return_value=FakePayload(
+        cros_payload.MAJOR_PAYLOAD_VERSION_CHROMEOS))
 
     with self.OutputCapturer() as output:
       payload_cmd.Run()
 
     stdout = output.GetStdout()
-    expected_out = """Payload version:         111
+    expected_out = """Payload version:         1
 Manifest length:         222
 Number of operations:    1
 Number of kernel ops:    1
@@ -152,17 +165,18 @@ Minor version:           4
 """
     self.assertEquals(stdout, expected_out)
 
-  def testListOps(self):
+  def testListOpsOnVersion1(self):
     """Verify that the --list_ops option gives the correct output."""
     payload_cmd = cros_payload.PayloadCommand(FakeOption(list_ops=True,
                                                          action='show'))
-    self.PatchObject(update_payload, 'Payload', return_value=FakePayload())
+    self.PatchObject(update_payload, 'Payload', return_value=FakePayload(
+        cros_payload.MAJOR_PAYLOAD_VERSION_CHROMEOS))
 
     with self.OutputCapturer() as output:
       payload_cmd.Run()
 
     stdout = output.GetStdout()
-    expected_out = """Payload version:         111
+    expected_out = """Payload version:         1
 Manifest length:         222
 Number of operations:    1
 Number of kernel ops:    1
@@ -185,17 +199,53 @@ Kernel install operations:
 """
     self.assertEquals(stdout, expected_out)
 
-  def testStats(self):
+  def testListOpsOnVersion2(self):
+    """Verify that the --list_ops option gives the correct output."""
+    payload_cmd = cros_payload.PayloadCommand(FakeOption(list_ops=True,
+                                                         action='show'))
+    self.PatchObject(update_payload, 'Payload', return_value=FakePayload(
+        cros_payload.MAJOR_PAYLOAD_VERSION_BRILLO))
+
+    with self.OutputCapturer() as output:
+      payload_cmd.Run()
+
+    stdout = output.GetStdout()
+    expected_out = """Payload version:         2
+Manifest length:         222
+Number of partitions:    2
+  Number of "rootfs" ops: 1
+  Number of "kernel" ops: 1
+Block size:              4096
+Minor version:           4
+
+rootfs install operations:
+  0: REPLACE_BZ
+    Data offset: 1
+    Data length: 1
+    Destination: 2 extents (3 blocks)
+      (1,1) (2,2)
+kernel install operations:
+  0: SOURCE_COPY
+    Source: 1 extent (1 block)
+      (1,1)
+    Destination: 20 extents (190 blocks)
+      (0,0) (1,1) (2,2) (3,3) (4,4) (5,5) (6,6) (7,7) (8,8) (9,9) (10,10)
+      (11,11) (12,12) (13,13) (14,14) (15,15) (16,16) (17,17) (18,18) (19,19)
+"""
+    self.assertEquals(stdout, expected_out)
+
+  def testStatsOnVersion1(self):
     """Verify that the --stats option works correctly."""
     payload_cmd = cros_payload.PayloadCommand(FakeOption(stats=True,
                                                          action='show'))
-    self.PatchObject(update_payload, 'Payload', return_value=FakePayload())
+    self.PatchObject(update_payload, 'Payload', return_value=FakePayload(
+        cros_payload.MAJOR_PAYLOAD_VERSION_CHROMEOS))
 
     with self.OutputCapturer() as output:
       payload_cmd.Run()
 
     stdout = output.GetStdout()
-    expected_out = """Payload version:         111
+    expected_out = """Payload version:         1
 Manifest length:         222
 Number of operations:    1
 Number of kernel ops:    1
@@ -207,17 +257,42 @@ Seeks when writing:      18
 """
     self.assertEquals(stdout, expected_out)
 
+  def testStatsOnVersion2(self):
+    """Verify that the --stats option works correctly on version 2."""
+    payload_cmd = cros_payload.PayloadCommand(FakeOption(stats=True,
+                                                         action='show'))
+    self.PatchObject(update_payload, 'Payload', return_value=FakePayload(
+        cros_payload.MAJOR_PAYLOAD_VERSION_BRILLO))
+
+    with self.OutputCapturer() as output:
+      payload_cmd.Run()
+
+    stdout = output.GetStdout()
+    expected_out = """Payload version:         2
+Manifest length:         222
+Number of partitions:    2
+  Number of "rootfs" ops: 1
+  Number of "kernel" ops: 1
+Block size:              4096
+Minor version:           4
+Blocks read:             11
+Blocks written:          193
+Seeks when writing:      18
+"""
+    self.assertEquals(stdout, expected_out)
+
   def testEmptySignatures(self):
     """Verify that the --signatures option works with unsigned payloads."""
     payload_cmd = cros_payload.PayloadCommand(
         FakeOption(action='show', signatures=True))
-    self.PatchObject(update_payload, 'Payload', return_value=FakePayload())
+    self.PatchObject(update_payload, 'Payload', return_value=FakePayload(
+        cros_payload.MAJOR_PAYLOAD_VERSION_CHROMEOS))
 
     with self.OutputCapturer() as output:
       payload_cmd.Run()
 
     stdout = output.GetStdout()
-    expected_out = """Payload version:         111
+    expected_out = """Payload version:         1
 Manifest length:         222
 Number of operations:    1
 Number of kernel ops:    1
@@ -232,7 +307,7 @@ No signatures stored in the payload
     """Verify that the --signatures option shows the present signatures."""
     payload_cmd = cros_payload.PayloadCommand(
         FakeOption(action='show', signatures=True))
-    payload = FakePayload()
+    payload = FakePayload(cros_payload.MAJOR_PAYLOAD_VERSION_CHROMEOS)
     payload.AddSignature(version=1, data='12345678abcdefgh\x00\x01\x02\x03')
     payload.AddSignature(data='I am a signature so access is yes.')
     self.PatchObject(update_payload, 'Payload', return_value=payload)
@@ -241,7 +316,7 @@ No signatures stored in the payload
       payload_cmd.Run()
 
     stdout = output.GetStdout()
-    expected_out = """Payload version:         111
+    expected_out = """Payload version:         1
 Manifest length:         222
 Number of operations:    1
 Number of kernel ops:    1