cros_generate_update_payload: Convert to chromite

This is a straightforward translation of the previous Bash
cros_generate_update_payload script. There are a number of low-hanging
fruit in terms of cleaning it up, but those will be the target of a
future CL.

BUG=chromium:860072
TEST=output matches old cros_generate_update_payload on same inputs;
  cros flash works

Change-Id: I2bf38a158d280e7f2d06cbed79ec5802415c7507
Reviewed-on: https://chromium-review.googlesource.com/1145683
Commit-Ready: ChromeOS CL Exonerator Bot <chromiumos-cl-exonerator@appspot.gserviceaccount.com>
Tested-by: Tudor Brindus <tbrindus@chromium.org>
Reviewed-by: Amin Hassani <ahassani@chromium.org>
Reviewed-by: Don Garrett <dgarrett@chromium.org>
Reviewed-by: Mike Frysinger <vapier@chromium.org>
diff --git a/scripts/cros_generate_update_payload_unittest.py b/scripts/cros_generate_update_payload_unittest.py
new file mode 100644
index 0000000..ddcb6c9
--- /dev/null
+++ b/scripts/cros_generate_update_payload_unittest.py
@@ -0,0 +1,126 @@
+# -*- coding: utf-8 -*-
+# Copyright 2018 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Unit tests for cros_generate_update_payload."""
+
+from __future__ import print_function
+
+import os
+
+from chromite.lib import cros_build_lib
+from chromite.lib import cros_test_lib
+from chromite.lib import osutils
+from chromite.scripts import cros_generate_update_payload
+
+
+class CopyFileSegmentTest(cros_test_lib.TempDirTestCase):
+  """Test cros_generate_update_payload.CopyFileSegment"""
+  def testCopyFileSegment(self):
+    """Test copying on a simple file."""
+    a = os.path.join(self.tempdir, 'a.txt')
+    osutils.WriteFile(a, '789')
+    b = os.path.join(self.tempdir, 'b.txt')
+    osutils.WriteFile(b, '0123')
+    cros_generate_update_payload.CopyFileSegment(a, 'r', 2, b, 'r+', in_seek=1)
+    self.assertEqual(osutils.ReadFile(b), '8923')
+
+
+class ExtractRootTest(cros_test_lib.MockTempDirTestCase):
+  """Test cros_generate_update_payload.ExtractRoot"""
+  def testTruncate(self):
+    """Test truncating on extraction."""
+    root = os.path.join(self.tempdir, 'root.bin')
+    root_pretruncate = os.path.join(self.tempdir, 'root_pretruncate.bin')
+
+    content = '0123456789'
+    osutils.WriteFile(root, content)
+    self.PatchObject(cros_generate_update_payload, 'ExtractPartitionToTempFile',
+                     return_value=root)
+    self.PatchObject(cros_generate_update_payload, 'Ext2FileSystemSize',
+                     return_value=2)
+
+    cros_generate_update_payload.ExtractRoot(None, root, root_pretruncate)
+    self.assertEqual(osutils.ReadFile(root), content[:2])
+    self.assertEqual(osutils.ReadFile(root_pretruncate), content)
+
+
+class Ext2FileSystemSizeTest(cros_test_lib.MockTestCase):
+  """Test cros_generate_update_payload.Ext2FileSystemSize"""
+  def testExt2FileSystemSize(self):
+    """Test on simple output."""
+    block_size = 4096
+    block_count = 123
+
+    self.PatchObject(cros_build_lib, 'RunCommand',
+                     return_value=cros_build_lib.CommandResult(output='''
+Block size: %d
+Other thing: 123456798
+Not an integer: cdsad132csda
+Block count: %d
+''' % (block_size, block_count)))
+
+    size = cros_generate_update_payload.Ext2FileSystemSize('/dev/null')
+    self.assertEqual(size, block_size * block_count)
+
+
+class ExtractPartitionToTempFileTest(cros_test_lib.MockTempDirTestCase):
+  """Test cros_generate_update_payload.ExtractPartitionToTempFile"""
+  def testExtractPartitionToTempFile(self):
+    """Tests extraction on a simple image."""
+    part_a_bin = '0123'
+    part_b_bin = '4567'
+    image_bin = part_a_bin + part_b_bin
+
+    image = os.path.join(self.tempdir, 'image.bin')
+    osutils.WriteFile(image, image_bin)
+    part_a = os.path.join(self.tempdir, 'a.bin')
+
+    fake_partitions = {
+        'PART-A': cros_build_lib.PartitionInfo(1, 0, 4, 4, 'fs', 'PART-A', ''),
+        'PART-B': cros_build_lib.PartitionInfo(2, 4, 8, 4, 'fs', 'PART-B', ''),
+    }
+    self.PatchObject(cros_build_lib, 'GetImageDiskPartitionInfo',
+                     return_value=fake_partitions)
+
+    cros_generate_update_payload.ExtractPartitionToTempFile(image, 'PART-A',
+                                                            temp_file=part_a)
+    self.assertEqual(osutils.ReadFile(part_a), part_a_bin)
+
+    # Make sure we correctly generate new temp files.
+    tmp = cros_generate_update_payload.ExtractPartitionToTempFile(image,
+                                                                  'PART-B')
+    self.assertEqual(osutils.ReadFile(tmp), part_b_bin)
+
+
+class DeltaGeneratorTest(cros_test_lib.RunCommandTempDirTestCase):
+  """Test correct arguments passed to delta_generator."""
+  def testDeltaGenerator(self):
+    """Test correct arguments propagated to delta_generator call."""
+    temp = os.path.join(self.tempdir, 'temp.bin')
+    osutils.WriteFile(temp, '0123456789')
+    self.PatchObject(cros_generate_update_payload, 'ExtractPartitionToTempFile',
+                     return_value=temp)
+    self.PatchObject(cros_generate_update_payload, 'Ext2FileSystemSize',
+                     return_value=4)
+
+    fake_partitions = {
+        'ROOT-A': cros_build_lib.PartitionInfo(3, 0, 4, 4, 'fs', 'ROOT-A', ''),
+    }
+    self.PatchObject(cros_build_lib, 'GetImageDiskPartitionInfo',
+                     return_value=fake_partitions)
+    cros_generate_update_payload.main([
+        '--image', '/dev/null',
+        '--src_image', '/dev/null',
+        '--output', '/dev/null',
+    ])
+
+    self.assertCommandContains([
+        '--major_version=1',
+        '--new_image=' + temp,
+        '--new_kernel=' + temp,
+        '--old_image=' + temp,
+        '--old_kernel=' + temp,
+        '--rootfs_partition_size=4',
+    ])