releasetools: Add tests for common.ReadApkCerts().

Test: python -m unittest test_common
Test: Run sign_target_files_apks.py on a target with compressed APKs.
Change-Id: I107a8b8f2f0f82e2d1947f14c8a8b3778f633b11
diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index 03e808f..ebebd63 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -792,11 +792,22 @@
 
 
 def ReadApkCerts(tf_zip):
-  """Given a target_files ZipFile, parse the META/apkcerts.txt file
-  and return a tuple with the following elements: (1) a dictionary that maps
-  packages to certs (based on the "certificate" and "private_key" attributes
-  in the file. (2) A string representing the extension of compressed APKs in
-  the target files (e.g ".gz" ".bro")."""
+  """Parses the APK certs info from a given target-files zip.
+
+  Given a target-files ZipFile, parses the META/apkcerts.txt entry and returns a
+  tuple with the following elements: (1) a dictionary that maps packages to
+  certs (based on the "certificate" and "private_key" attributes in the file;
+  (2) a string representing the extension of compressed APKs in the target files
+  (e.g ".gz", ".bro").
+
+  Args:
+    tf_zip: The input target_files ZipFile (already open).
+
+  Returns:
+    (certmap, ext): certmap is a dictionary that maps packages to certs; ext is
+        the extension string of compressed APKs (e.g. ".gz"), or None if there's
+        no compressed APKs.
+  """
   certmap = {}
   compressed_extension = None
 
@@ -812,41 +823,51 @@
     line = line.strip()
     if not line:
       continue
-    m = re.match(r'^name="(?P<NAME>.*)"\s+certificate="(?P<CERT>.*)"\s+'
-                 r'private_key="(?P<PRIVKEY>.*?)"(\s+compressed="(?P<COMPRESSED>.*)")?$',
-                 line)
-    if m:
-      matches = m.groupdict()
-      cert = matches["CERT"]
-      privkey = matches["PRIVKEY"]
-      name = matches["NAME"]
-      this_compressed_extension = matches["COMPRESSED"]
-      public_key_suffix_len = len(OPTIONS.public_key_suffix)
-      private_key_suffix_len = len(OPTIONS.private_key_suffix)
-      if cert in SPECIAL_CERT_STRINGS and not privkey:
-        certmap[name] = cert
-      elif (cert.endswith(OPTIONS.public_key_suffix) and
-            privkey.endswith(OPTIONS.private_key_suffix) and
-            cert[:-public_key_suffix_len] == privkey[:-private_key_suffix_len]):
-        certmap[name] = cert[:-public_key_suffix_len]
-      else:
-        raise ValueError("failed to parse line from apkcerts.txt:\n" + line)
-      if this_compressed_extension:
-        # Only count the installed files.
-        filename = name + '.' + this_compressed_extension
-        if filename not in installed_files:
-          continue
-        # Make sure that all the values in the compression map have the same
-        # extension. We don't support multiple compression methods in the same
-        # system image.
-        if compressed_extension:
-          if this_compressed_extension != compressed_extension:
-            raise ValueError("multiple compressed extensions : %s vs %s",
-                             (compressed_extension, this_compressed_extension))
-        else:
-          compressed_extension = this_compressed_extension
+    m = re.match(
+        r'^name="(?P<NAME>.*)"\s+certificate="(?P<CERT>.*)"\s+'
+        r'private_key="(?P<PRIVKEY>.*?)"(\s+compressed="(?P<COMPRESSED>.*)")?$',
+        line)
+    if not m:
+      continue
 
-  return (certmap, ("." + compressed_extension) if compressed_extension else None)
+    matches = m.groupdict()
+    cert = matches["CERT"]
+    privkey = matches["PRIVKEY"]
+    name = matches["NAME"]
+    this_compressed_extension = matches["COMPRESSED"]
+
+    public_key_suffix_len = len(OPTIONS.public_key_suffix)
+    private_key_suffix_len = len(OPTIONS.private_key_suffix)
+    if cert in SPECIAL_CERT_STRINGS and not privkey:
+      certmap[name] = cert
+    elif (cert.endswith(OPTIONS.public_key_suffix) and
+          privkey.endswith(OPTIONS.private_key_suffix) and
+          cert[:-public_key_suffix_len] == privkey[:-private_key_suffix_len]):
+      certmap[name] = cert[:-public_key_suffix_len]
+    else:
+      raise ValueError("Failed to parse line from apkcerts.txt:\n" + line)
+
+    if not this_compressed_extension:
+      continue
+
+    # Only count the installed files.
+    filename = name + '.' + this_compressed_extension
+    if filename not in installed_files:
+      continue
+
+    # Make sure that all the values in the compression map have the same
+    # extension. We don't support multiple compression methods in the same
+    # system image.
+    if compressed_extension:
+      if this_compressed_extension != compressed_extension:
+        raise ValueError(
+            "Multiple compressed extensions: {} vs {}".format(
+                compressed_extension, this_compressed_extension))
+    else:
+      compressed_extension = this_compressed_extension
+
+  return (certmap,
+          ("." + compressed_extension) if compressed_extension else None)
 
 
 COMMON_DOCSTRING = """
diff --git a/tools/releasetools/test_common.py b/tools/releasetools/test_common.py
index ed454ca..8fb4600 100644
--- a/tools/releasetools/test_common.py
+++ b/tools/releasetools/test_common.py
@@ -353,6 +353,128 @@
       os.remove(zip_file.name)
 
 
+class CommonApkUtilsTest(unittest.TestCase):
+  """Tests the APK utils related functions."""
+
+  APKCERTS_TXT1 = (
+      'name="RecoveryLocalizer.apk" certificate="certs/devkey.x509.pem"'
+      ' private_key="certs/devkey.pk8"\n'
+      'name="Settings.apk"'
+      ' certificate="build/target/product/security/platform.x509.pem"'
+      ' private_key="build/target/product/security/platform.pk8"\n'
+      'name="TV.apk" certificate="PRESIGNED" private_key=""\n'
+  )
+
+  APKCERTS_CERTMAP1 = {
+      'RecoveryLocalizer.apk' : 'certs/devkey',
+      'Settings.apk' : 'build/target/product/security/platform',
+      'TV.apk' : 'PRESIGNED',
+  }
+
+  APKCERTS_TXT2 = (
+      'name="Compressed1.apk" certificate="certs/compressed1.x509.pem"'
+      ' private_key="certs/compressed1.pk8" compressed="gz"\n'
+      'name="Compressed2a.apk" certificate="certs/compressed2.x509.pem"'
+      ' private_key="certs/compressed2.pk8" compressed="gz"\n'
+      'name="Compressed2b.apk" certificate="certs/compressed2.x509.pem"'
+      ' private_key="certs/compressed2.pk8" compressed="gz"\n'
+      'name="Compressed3.apk" certificate="certs/compressed3.x509.pem"'
+      ' private_key="certs/compressed3.pk8" compressed="gz"\n'
+  )
+
+  APKCERTS_CERTMAP2 = {
+      'Compressed1.apk' : 'certs/compressed1',
+      'Compressed2a.apk' : 'certs/compressed2',
+      'Compressed2b.apk' : 'certs/compressed2',
+      'Compressed3.apk' : 'certs/compressed3',
+  }
+
+  APKCERTS_TXT3 = (
+      'name="Compressed4.apk" certificate="certs/compressed4.x509.pem"'
+      ' private_key="certs/compressed4.pk8" compressed="xz"\n'
+  )
+
+  APKCERTS_CERTMAP3 = {
+      'Compressed4.apk' : 'certs/compressed4',
+  }
+
+  def tearDown(self):
+    common.Cleanup()
+
+  @staticmethod
+  def _write_apkcerts_txt(apkcerts_txt, additional=None):
+    if additional is None:
+      additional = []
+    target_files = common.MakeTempFile(suffix='.zip')
+    with zipfile.ZipFile(target_files, 'w') as target_files_zip:
+      target_files_zip.writestr('META/apkcerts.txt', apkcerts_txt)
+      for entry in additional:
+        target_files_zip.writestr(entry, '')
+    return target_files
+
+  def test_ReadApkCerts_NoncompressedApks(self):
+    target_files = self._write_apkcerts_txt(self.APKCERTS_TXT1)
+    with zipfile.ZipFile(target_files, 'r') as input_zip:
+      certmap, ext = common.ReadApkCerts(input_zip)
+
+    self.assertDictEqual(self.APKCERTS_CERTMAP1, certmap)
+    self.assertIsNone(ext)
+
+  def test_ReadApkCerts_CompressedApks(self):
+    # We have "installed" Compressed1.apk.gz only. Note that Compressed3.apk is
+    # not stored in '.gz' format, so it shouldn't be considered as installed.
+    target_files = self._write_apkcerts_txt(
+        self.APKCERTS_TXT2,
+        ['Compressed1.apk.gz', 'Compressed3.apk'])
+
+    with zipfile.ZipFile(target_files, 'r') as input_zip:
+      certmap, ext = common.ReadApkCerts(input_zip)
+
+    self.assertDictEqual(self.APKCERTS_CERTMAP2, certmap)
+    self.assertEqual('.gz', ext)
+
+    # Alternative case with '.xz'.
+    target_files = self._write_apkcerts_txt(
+        self.APKCERTS_TXT3, ['Compressed4.apk.xz'])
+
+    with zipfile.ZipFile(target_files, 'r') as input_zip:
+      certmap, ext = common.ReadApkCerts(input_zip)
+
+    self.assertDictEqual(self.APKCERTS_CERTMAP3, certmap)
+    self.assertEqual('.xz', ext)
+
+  def test_ReadApkCerts_CompressedAndNoncompressedApks(self):
+    target_files = self._write_apkcerts_txt(
+        self.APKCERTS_TXT1 + self.APKCERTS_TXT2,
+        ['Compressed1.apk.gz', 'Compressed3.apk'])
+
+    with zipfile.ZipFile(target_files, 'r') as input_zip:
+      certmap, ext = common.ReadApkCerts(input_zip)
+
+    certmap_merged = self.APKCERTS_CERTMAP1.copy()
+    certmap_merged.update(self.APKCERTS_CERTMAP2)
+    self.assertDictEqual(certmap_merged, certmap)
+    self.assertEqual('.gz', ext)
+
+  def test_ReadApkCerts_MultipleCompressionMethods(self):
+    target_files = self._write_apkcerts_txt(
+        self.APKCERTS_TXT2 + self.APKCERTS_TXT3,
+        ['Compressed1.apk.gz', 'Compressed4.apk.xz'])
+
+    with zipfile.ZipFile(target_files, 'r') as input_zip:
+      self.assertRaises(ValueError, common.ReadApkCerts, input_zip)
+
+  def test_ReadApkCerts_MismatchingKeys(self):
+    malformed_apkcerts_txt = (
+        'name="App1.apk" certificate="certs/cert1.x509.pem"'
+        ' private_key="certs/cert2.pk8"\n'
+    )
+    target_files = self._write_apkcerts_txt(malformed_apkcerts_txt)
+
+    with zipfile.ZipFile(target_files, 'r') as input_zip:
+      self.assertRaises(ValueError, common.ReadApkCerts, input_zip)
+
+
 class InstallRecoveryScriptFormatTest(unittest.TestCase):
   """Checks the format of install-recovery.sh.