update payload : Add unittest for paycheck.py

Use generated sample payloads to verify and apply each payload.

BUG=chromium:1028646
TEST=sudo FEATURES=test emerge update_payload
TEST=./generate_payloads

Cq-Depend: chromium:2401388
Change-Id: I2b817c4b71edf4cc6bd36d9ee021366818a42ebb
Reviewed-on: https://chromium-review.googlesource.com/c/aosp/platform/system/update_engine/+/2401389
Tested-by: Vyshu Khota <vyshu@google.com>
Commit-Queue: Vyshu Khota <vyshu@google.com>
Reviewed-by: Amin Hassani <ahassani@chromium.org>
Reviewed-by: Jae Hoon Kim <kimjae@chromium.org>
diff --git a/scripts/paycheck_unittest.py b/scripts/paycheck_unittest.py
new file mode 100755
index 0000000..e54a3c0
--- /dev/null
+++ b/scripts/paycheck_unittest.py
@@ -0,0 +1,105 @@
+#!/usr/bin/env python
+#
+# Copyright (C) 2020 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Unit testing paycheck.py."""
+
+# This test requires new (Y) and old (X) images, as well as a full payload
+# from image Y and a delta payload from Y to X for each partition.
+# Payloads are from sample_images/generate_payloads.
+#
+# The test performs the following:
+#
+# - It statically applies the full and delta payloads.
+#
+# - It applies full_payload to yield a new kernel (kern.part) and rootfs
+#   (root.part) and compares them to the new image partitions.
+#
+# - It applies delta_payload to the old image to yield a new kernel and rootfs
+#   and compares them to the new image partitions.
+#
+# Previously test_paycheck.sh. Run with update_payload ebuild.
+
+# Disable check for function names to avoid errors based on old code
+# pylint: disable-msg=invalid-name
+
+import filecmp
+import os
+import subprocess
+import unittest
+
+
+class PaycheckTest(unittest.TestCase):
+  """Test paycheck functions."""
+
+  def setUp(self):
+    self.tmpdir = os.getenv('T')
+
+    self._full_payload = os.path.join(self.tmpdir, 'full_payload.bin')
+    self._delta_payload = os.path.join(self.tmpdir, 'delta_payload.bin')
+
+    self._new_kernel = os.path.join(self.tmpdir, 'disk_ext2_4k.img')
+    self._new_root = os.path.join(self.tmpdir, 'disk_sqfs_default.img')
+    self._old_kernel = os.path.join(self.tmpdir,
+                                    'disk_ext2_4k_empty.img')
+    self._old_root = os.path.join(self.tmpdir, 'disk_sqfs_empty.img')
+
+    # Temp output files.
+    self._kernel_part = os.path.join(self.tmpdir, 'kern.part')
+    self._root_part = os.path.join(self.tmpdir, 'root.part')
+
+  def checkPayload(self, type_arg, payload):
+    """Checks Payload."""
+    self.assertEqual(0, subprocess.check_call(['./paycheck.py', '-t',
+                                               type_arg, payload]))
+
+  def testFullPayload(self):
+    """Checks the full payload statically."""
+    self.checkPayload('full', self._full_payload)
+
+  def testDeltaPayload(self):
+    """Checks the delta payload statically."""
+    self.checkPayload('delta', self._delta_payload)
+
+  def testApplyFullPayload(self):
+    """Applies full payloads and compares results to new sample images."""
+    self.assertEqual(0, subprocess.check_call(['./paycheck.py',
+                                               self._full_payload,
+                                               '--part_names', 'kernel', 'root',
+                                               '--out_dst_part_paths',
+                                               self._kernel_part,
+                                               self._root_part]))
+
+    # Check if generated full image is equal to sample image.
+    self.assertTrue(filecmp.cmp(self._kernel_part, self._new_kernel))
+    self.assertTrue(filecmp.cmp(self._root_part, self._new_root))
+
+  def testApplyDeltaPayload(self):
+    """Applies delta to old image and checks against new sample images."""
+    self.assertEqual(0, subprocess.check_call(['./paycheck.py',
+                                               self._delta_payload,
+                                               '--part_names', 'kernel', 'root',
+                                               '--src_part_paths',
+                                               self._old_kernel, self._old_root,
+                                               '--out_dst_part_paths',
+                                               self._kernel_part,
+                                               self._root_part]))
+
+    self.assertTrue(filecmp.cmp(self._kernel_part, self._new_kernel))
+    self.assertTrue(filecmp.cmp(self._root_part, self._new_root))
+
+if __name__ == '__main__':
+  unittest.main()
diff --git a/scripts/run_unittests b/scripts/run_unittests
index 0d301ba..db5ed73 100755
--- a/scripts/run_unittests
+++ b/scripts/run_unittests
@@ -26,5 +26,6 @@
 done
 
 ./payload_info_unittest.py
+./paycheck_unittest.py
 
 exit 0
diff --git a/scripts/test_paycheck.sh b/scripts/test_paycheck.sh
deleted file mode 100755
index 239b984..0000000
--- a/scripts/test_paycheck.sh
+++ /dev/null
@@ -1,169 +0,0 @@
-#!/bin/bash
-#
-# Copyright (C) 2013 The Android Open Source Project
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
-# A test script for paycheck.py and the update_payload.py library.
-#
-# This script requires three payload files, along with a metadata signature for
-# each, and a public key for verifying signatures. Payload include:
-#
-# - A full payload for release X (old_full_payload)
-#
-# - A full payload for release Y (new_full_payload), where Y > X
-#
-# - A delta payload from X to Y (delta_payload)
-#
-# The test performs the following:
-#
-# - It verifies each payload against its metadata signature, also asserting the
-#   payload type. Another artifact is a human-readable payload report, which
-#   is output to stdout to be inspected by the user.
-#
-# - It applies old_full_payload to yield old kernel (old_kern.part) and rootfs
-#   (old_root.part) partitions.
-#
-# - It applies delta_payload to old_{kern,root}.part to yield new kernel
-#   (new_delta_kern.part) and rootfs (new_delta_root.part) partitions.
-#
-# - It applies new_full_payload to yield reference new kernel
-#   (new_full_kern.part) and rootfs (new_full_root.part) partitions.
-#
-# - It compares new_{delta,full}_kern.part and new_{delta,full}_root.part to
-#   ensure that they are binary identical.
-#
-# If all steps have completed successfully we know with high certainty that
-# paycheck.py (and hence update_payload.py) correctly parses both full and delta
-# payloads, and applies them to yield the expected result. Finally, each
-# paycheck.py execution is timed.
-
-
-# Stop on errors, unset variables.
-set -e
-set -u
-
-# Temporary image files.
-OLD_KERN_PART=old_kern.part
-OLD_ROOT_PART=old_root.part
-NEW_DELTA_KERN_PART=new_delta_kern.part
-NEW_DELTA_ROOT_PART=new_delta_root.part
-NEW_FULL_KERN_PART=new_full_kern.part
-NEW_FULL_ROOT_PART=new_full_root.part
-CROS_PARTS="kernel root"
-
-
-log() {
-  echo "$@" >&2
-}
-
-die() {
-  log "$@"
-  exit 1
-}
-
-usage_and_exit() {
-  cat >&2 <<EOF
-Usage: ${0##*/} old_full_payload delta_payload new_full_payload
-EOF
-  exit
-}
-
-check_payload() {
-  payload_file=$1
-  payload_type=$2
-
-  time ${paycheck} -t ${payload_type} ${payload_file}
-}
-
-apply_full_payload() {
-  payload_file=$1
-  out_dst_kern_part="$2/$3"
-  out_dst_root_part="$2/$4"
-
-  time ${paycheck} ${payload_file} \
-    --part_names ${CROS_PARTS} \
-    --out_dst_part_paths ${out_dst_kern_part} ${out_dst_root_part}
-}
-
-apply_delta_payload() {
-  payload_file=$1
-  out_dst_kern_part="$2/$3"
-  out_dst_root_part="$2/$4"
-  dst_kern_part="$2/$5"
-  dst_root_part="$2/$6"
-  src_kern_part="$2/$7"
-  src_root_part="$2/$8"
-
-  time ${paycheck} ${payload_file} \
-    --part_names ${CROS_PARTS} \
-    --out_dst_part_paths ${out_dst_kern_part} ${out_dst_root_part} \
-    --dst_part_paths ${dst_kern_part} ${dst_root_part} \
-    --src_part_paths ${src_kern_part} ${src_root_part}
-}
-
-main() {
-  # Read command-line arguments.
-  if [ $# == 1 ] && [ "$1" == "-h" ]; then
-    usage_and_exit
-  elif [ $# != 3 ]; then
-    die "Error: unexpected number of arguments"
-  fi
-  old_full_payload="$1"
-  delta_payload="$2"
-  new_full_payload="$3"
-
-  # Find paycheck.py
-  paycheck=${0%/*}/paycheck.py
-  if [ -z "${paycheck}" ] || [ ! -x ${paycheck} ]; then
-    die "cannot find ${paycheck} or file is not executable"
-  fi
-
-  # Check the payloads statically.
-  log "Checking payloads..."
-  check_payload "${old_full_payload}" full
-  check_payload "${new_full_payload}" full
-  check_payload "${delta_payload}" delta
-  log "Done"
-
-  # Apply full/delta payloads and verify results are identical.
-  tmpdir="$(mktemp -d --tmpdir test_paycheck.XXXXXXXX)"
-  log "Initiating application of payloads at $tmpdir"
-
-  log "Applying old full payload..."
-  apply_full_payload "${old_full_payload}" "${tmpdir}" "${OLD_KERN_PART}" \
-    "${OLD_ROOT_PART}"
-  log "Done"
-
-  log "Applying new full payload..."
-  apply_full_payload "${new_full_payload}" "${tmpdir}" "${NEW_FULL_KERN_PART}" \
-    "${NEW_FULL_ROOT_PART}"
-  log "Done"
-
-  log "Applying delta payload to old partitions..."
-  apply_delta_payload "${delta_payload}" "${tmpdir}" "${NEW_DELTA_KERN_PART}" \
-    "${NEW_DELTA_ROOT_PART}" "${NEW_FULL_KERN_PART}" \
-    "${NEW_FULL_ROOT_PART}" "${OLD_KERN_PART}" "${OLD_ROOT_PART}"
-  log "Done"
-
-  log "Comparing results of delta and new full updates..."
-  diff "${tmpdir}/${NEW_FULL_KERN_PART}" "${tmpdir}/${NEW_DELTA_KERN_PART}"
-  diff "${tmpdir}/${NEW_FULL_ROOT_PART}" "${tmpdir}/${NEW_DELTA_ROOT_PART}"
-  log "Done"
-
-  log "Cleaning up"
-  rm -fr "${tmpdir}"
-}
-
-main "$@"