Merge "releasetools: Stop copying images from RADIO/ to IMAGES/."

tools/releasetools/add_img_to_target_files.py

@@ -485,18 +485,15 @@ def AddCache(output_zip):
   img.Write()
 
 
-def AddRadioImagesForAbOta(output_zip, ab_partitions):
-  """Adds the radio images needed for A/B OTA to the output file.
+def CheckAbOtaImages(output_zip, ab_partitions):
+  """Checks that all the listed A/B partitions have their images available.
 
-  It parses the list of A/B partitions, looks for the missing ones from RADIO/,
-  and copies them to IMAGES/ of the output file (or dir).
-
-  It also ensures that on returning from the function all the listed A/B
-  partitions must have their images available under IMAGES/.
+  The images need to be available under IMAGES/ or RADIO/, with the former takes
+  a priority.
 
   Args:
     output_zip: The output zip file (needs to be already open), or None to
-        write images to OPTIONS.input_tmp/.
+        find images in OPTIONS.input_tmp/.
     ab_partitions: The list of A/B partitions.
 
   Raises:
@@ -504,27 +501,20 @@ def AddRadioImagesForAbOta(output_zip, ab_partitions):
   """
   for partition in ab_partitions:
     img_name = partition.strip() + ".img"
-    prebuilt_path = os.path.join(OPTIONS.input_tmp, "IMAGES", img_name)
-    if os.path.exists(prebuilt_path):
-      print("%s already exists, no need to overwrite..." % (img_name,))
-      continue
-
-    img_radio_path = os.path.join(OPTIONS.input_tmp, "RADIO", img_name)
-    if os.path.exists(img_radio_path):
-      if output_zip:
-        common.ZipWrite(output_zip, img_radio_path, "IMAGES/" + img_name)
-      else:
-        shutil.copy(img_radio_path, prebuilt_path)
-      continue
 
     # Assert that the image is present under IMAGES/ now.
     if output_zip:
       # Zip spec says: All slashes MUST be forward slashes.
-      img_path = 'IMAGES/' + img_name
-      assert img_path in output_zip.namelist(), "cannot find " + img_name
+      images_path = "IMAGES/" + img_name
+      radio_path = "RADIO/" + img_name
+      available = (images_path in output_zip.namelist() or
+                   radio_path in output_zip.namelist())
     else:
-      img_path = os.path.join(OPTIONS.input_tmp, "IMAGES", img_name)
-      assert os.path.exists(img_path), "cannot find " + img_name
+      images_path = os.path.join(OPTIONS.input_tmp, "IMAGES", img_name)
+      radio_path = os.path.join(OPTIONS.input_tmp, "RADIO", img_name)
+      available = os.path.exists(images_path) or os.path.exists(radio_path)
+
+    assert available, "Failed to find " + img_name
 
 
 def AddCareMapTxtForAbOta(output_zip, ab_partitions, image_paths):
@@ -751,10 +741,9 @@ def AddImagesToTargetFiles(filename):
     with open(ab_partitions_txt, 'r') as f:
       ab_partitions = f.readlines()
 
-    # For devices using A/B update, copy over images from RADIO/ to IMAGES/ and
-    # make sure we have all the needed images ready under IMAGES/. All images
-    # should have '.img' as extension.
-    AddRadioImagesForAbOta(output_zip, ab_partitions)
+    # For devices using A/B update, make sure we have all the needed images
+    # ready under IMAGES/ or RADIO/.
+    CheckAbOtaImages(output_zip, ab_partitions)
 
     # Generate care_map.txt for system and vendor partitions (if present), then
     # write this file to target_files package.

tools/releasetools/test_add_img_to_target_files.py

@@ -22,8 +22,7 @@ import zipfile
 import common
 import test_utils
 from add_img_to_target_files import (
-    AddCareMapTxtForAbOta, AddPackRadioImages, AddRadioImagesForAbOta,
-    GetCareMap)
+    AddCareMapTxtForAbOta, AddPackRadioImages, CheckAbOtaImages, GetCareMap)
 from rangelib import RangeSet
 
 
@@ -55,49 +54,25 @@ class AddImagesToTargetFilesTest(unittest.TestCase):
       os.mkdir(images_path)
     return images, images_path
 
-  def test_AddRadioImagesForAbOta_imageExists(self):
+  def test_CheckAbOtaImages_imageExistsUnderImages(self):
     """Tests the case with existing images under IMAGES/."""
-    images, images_path = self._create_images(['aboot', 'xbl'], 'IMAGES')
-    AddRadioImagesForAbOta(None, images)
+    images, _ = self._create_images(['aboot', 'xbl'], 'IMAGES')
+    CheckAbOtaImages(None, images)
 
-    for image in images:
-      self.assertTrue(
-          os.path.exists(os.path.join(images_path, image + '.img')))
+  def test_CheckAbOtaImages_imageExistsUnderRadio(self):
+    """Tests the case with some image under RADIO/."""
+    images, _ = self._create_images(['system', 'vendor'], 'IMAGES')
+    radio_path = os.path.join(OPTIONS.input_tmp, 'RADIO')
+    if not os.path.exists(radio_path):
+      os.mkdir(radio_path)
+    with open(os.path.join(radio_path, 'modem.img'), 'wb') as image_fp:
+      image_fp.write('modem'.encode())
+    CheckAbOtaImages(None, images + ['modem'])
 
-  def test_AddRadioImagesForAbOta_copyFromRadio(self):
-    """Tests the case that copies images from RADIO/."""
-    images, images_path = self._create_images(['aboot', 'xbl'], 'RADIO')
-    AddRadioImagesForAbOta(None, images)
-
-    for image in images:
-      self.assertTrue(
-          os.path.exists(os.path.join(images_path, image + '.img')))
-
-  def test_AddRadioImagesForAbOta_copyFromRadio_zipOutput(self):
+  def test_CheckAbOtaImages_missingImages(self):
     images, _ = self._create_images(['aboot', 'xbl'], 'RADIO')
-
-    # Set up the output zip.
-    output_file = common.MakeTempFile(suffix='.zip')
-    with zipfile.ZipFile(output_file, 'w') as output_zip:
-      AddRadioImagesForAbOta(output_zip, images)
-
-    with zipfile.ZipFile(output_file, 'r') as verify_zip:
-      for image in images:
-        self.assertIn('IMAGES/' + image + '.img', verify_zip.namelist())
-
-  def test_AddRadioImagesForAbOta_missingImages(self):
-    images, _ = self._create_images(['aboot', 'xbl'], 'RADIO')
-    self.assertRaises(AssertionError, AddRadioImagesForAbOta, None,
-                      images + ['baz'])
-
-  def test_AddRadioImagesForAbOta_missingImages_zipOutput(self):
-    images, _ = self._create_images(['aboot', 'xbl'], 'RADIO')
-
-    # Set up the output zip.
-    output_file = common.MakeTempFile(suffix='.zip')
-    with zipfile.ZipFile(output_file, 'w') as output_zip:
-      self.assertRaises(AssertionError, AddRadioImagesForAbOta, output_zip,
-                        images + ['baz'])
+    self.assertRaises(
+        AssertionError, CheckAbOtaImages, None, images + ['baz'])
 
   def test_AddPackRadioImages(self):
     images, images_path = self._create_images(['foo', 'bar'], 'RADIO')

tools/releasetools/test_ota_from_target_files.py

@@ -50,17 +50,25 @@ def construct_target_files(secondary=False):
             "POSTINSTALL_OPTIONAL_system=true",
         ]))
 
+    ab_partitions = [
+        ('IMAGES', 'boot'),
+        ('IMAGES', 'system'),
+        ('IMAGES', 'vendor'),
+        ('RADIO', 'bootloader'),
+        ('RADIO', 'modem'),
+    ]
     # META/ab_partitions.txt
-    ab_partitions = ['boot', 'system', 'vendor']
     target_files_zip.writestr(
         'META/ab_partitions.txt',
-        '\n'.join(ab_partitions))
+        '\n'.join([partition[1] for partition in ab_partitions]))
 
     # Create dummy images for each of them.
-    for partition in ab_partitions:
-      target_files_zip.writestr('IMAGES/' + partition + '.img',
-                                os.urandom(len(partition)))
+    for path, partition in ab_partitions:
+      target_files_zip.writestr(
+          '{}/{}.img'.format(path, partition),
+          os.urandom(len(partition)))
 
+    # system_other shouldn't appear in META/ab_partitions.txt.
     if secondary:
       target_files_zip.writestr('IMAGES/system_other.img',
                                 os.urandom(len("system_other")))