releasetools: Fix the use of StringIO.

Based on the actual semantics, it actually wants an in-memory _bytes_
buffer (io.BytesIO), especially when running with Python 3. This CL
fixes the issue and adds a unittest.

Bug: 131631303
Test: python -m unittest test_sign_target_files_apks
Test: python3 -m unittest test_sign_target_files_apks
Change-Id: I3fb067acc26713f1842e831225607779fd0d1b7e

2 files changed, 20 insertions, 7 deletions

diff --git a/tools/releasetools/sign_target_files_apks.py b/tools/releasetools/sign_target_files_apks.py
index 4cb3a379f7..2f3642b227 100755
--- a/tools/releasetools/sign_target_files_apks.py
+++ b/tools/releasetools/sign_target_files_apks.py
@@ -111,6 +111,7 @@ import base64
 import copy
 import errno
 import gzip
+import io
 import itertools
 import logging
 import os
@@ -746,12 +747,7 @@ def WriteOtacerts(output_zip, filename, keys):
     filename: The archive name in the output zip.
     keys: A list of public keys to use during OTA package verification.
   """
-
-  try:
-    from StringIO import StringIO
-  except ImportError:
-    from io import StringIO
-  temp_file = StringIO()
+  temp_file = io.BytesIO()
   certs_zip = zipfile.ZipFile(temp_file, "w")
   for k in keys:
     common.ZipWrite(certs_zip, k)
diff --git a/tools/releasetools/test_sign_target_files_apks.py b/tools/releasetools/test_sign_target_files_apks.py
index 01007290ae..70c147e3f7 100644
--- a/tools/releasetools/test_sign_target_files_apks.py
+++ b/tools/releasetools/test_sign_target_files_apks.py
@@ -15,6 +15,7 @@
 #
 
 import base64
+import io
 import os.path
 import zipfile
 
@@ -22,7 +23,7 @@ import common
 import test_utils
 from sign_target_files_apks import (
     CheckApkAndApexKeysAvailable, EditTags, GetApkFileInfo, ReadApexKeysInfo,
-    ReplaceCerts, ReplaceVerityKeyId, RewriteProps)
+    ReplaceCerts, ReplaceVerityKeyId, RewriteProps, WriteOtacerts)
 
 
 class SignTargetFilesApksTest(test_utils.ReleaseToolsTestCase):
@@ -236,6 +237,22 @@ name="apex.apexd_test_different_app.apex" public_key="system/apex/apexd/apexd_te
     }
     self.assertEqual(output_xml, ReplaceCerts(input_xml))
 
+  def test_WriteOtacerts(self):
+    certs = [
+        os.path.join(self.testdata_dir, 'platform.x509.pem'),
+        os.path.join(self.testdata_dir, 'media.x509.pem'),
+        os.path.join(self.testdata_dir, 'testkey.x509.pem'),
+    ]
+    entry_name = 'SYSTEM/etc/security/otacerts.zip'
+    output_file = common.MakeTempFile(suffix='.zip')
+    with zipfile.ZipFile(output_file, 'w') as output_zip:
+      WriteOtacerts(output_zip, entry_name, certs)
+    with zipfile.ZipFile(output_file) as input_zip:
+      self.assertIn(entry_name, input_zip.namelist())
+      otacerts_file = io.BytesIO(input_zip.read(entry_name))
+      with zipfile.ZipFile(otacerts_file) as otacerts_zip:
+        self.assertEqual(3, len(otacerts_zip.namelist()))
+
   def test_CheckApkAndApexKeysAvailable(self):
     input_file = common.MakeTempFile(suffix='.zip')
     with zipfile.ZipFile(input_file, 'w') as input_zip: