From bb2658d5e2b9cbf3bf9ff30dc25f4f3aeb85e94a Mon Sep 17 00:00:00 2001
From: Yifan Hong <elsk@google.com>
Date: Fri, 25 Jan 2019 12:30:58 -0800
Subject: [PATCH] releasetools: Really allow removing partitions

- Add hashtree_info to EmptyImage so that BlockDifference.Compute()
can accept EmptyImage() as target image, which is the case when
a partition is removed.

- BlockDifference also checks source_info_dict to determine
whether a partition is dynamic. When a partition is removed,
its name does not appear in target_info_dict.

- Add tests to ensure DynamicPartitionDifference() still works.

Test: DynamicPartitionDifferenceTest
Test: test_blockimgdiff

Change-Id: Iadb1db075f5dc344db6d5ade358c83b01231e443
---
 tools/releasetools/blockimgdiff.py |  1 +
 tools/releasetools/common.py       | 12 +++++++---
 tools/releasetools/test_common.py  | 38 +++++++++++++++++++++++++++++-
 3 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/tools/releasetools/blockimgdiff.py b/tools/releasetools/blockimgdiff.py
index 456f791ea..e5a9050b1 100644
--- a/tools/releasetools/blockimgdiff.py
+++ b/tools/releasetools/blockimgdiff.py
@@ -84,6 +84,7 @@ class EmptyImage(Image):
     self.extended = RangeSet()
     self.total_blocks = 0
     self.file_map = {}
+    self.hashtree_info = None
 
   def RangeSha1(self, ranges):
     return sha1().hexdigest()
diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index 04158943d..5d310d2a1 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -1770,15 +1770,21 @@ class BlockDifference(object):
     if OPTIONS.source_info_dict is None:
       is_dynamic_build = OPTIONS.info_dict.get(
           "use_dynamic_partitions") == "true"
+      is_dynamic_source = False
     else:
       is_dynamic_build = OPTIONS.source_info_dict.get(
           "use_dynamic_partitions") == "true"
+      is_dynamic_source = partition in shlex.split(
+          OPTIONS.source_info_dict.get("dynamic_partition_list", "").strip())
 
-    # For dynamic partitions builds, always check partition list in target build
-    # because new partitions may be added.
-    is_dynamic = is_dynamic_build and partition in shlex.split(
+    is_dynamic_target = partition in shlex.split(
         OPTIONS.info_dict.get("dynamic_partition_list", "").strip())
 
+    # For dynamic partitions builds, check partition list in both source
+    # and target build because new partitions may be added, and existing
+    # partitions may be removed.
+    is_dynamic = is_dynamic_build and (is_dynamic_source or is_dynamic_target)
+
     if is_dynamic:
       self.device = 'map_partition("%s")' % partition
     else:
diff --git a/tools/releasetools/test_common.py b/tools/releasetools/test_common.py
index 361963316..cfd070d81 100644
--- a/tools/releasetools/test_common.py
+++ b/tools/releasetools/test_common.py
@@ -27,6 +27,7 @@ import test_utils
 import validate_target_files
 from rangelib import RangeSet
 
+from blockimgdiff import EmptyImage, DataImage
 
 KiB = 1024
 MiB = 1024 * KiB
@@ -1113,7 +1114,7 @@ super_group_qux_group_size={group_qux_size}
                     "ops that remove / shrink partitions must precede ops that "
                     "grow / add partitions")
 
-  def test_inc_partitions(self):
+  def test_incremental(self):
     source_info = common.LoadDictionaryFromLines("""
 dynamic_partition_list=system vendor product product_services
 super_partition_groups=group_foo
@@ -1186,3 +1187,38 @@ super_group_bar_partition_list=product
                     min_idx_move_partition_in_foo,
                     "Must shrink partitions / remove partitions from group"
                     "before adding / moving partitions into group")
+
+  def test_remove_partition(self):
+    source_info = common.LoadDictionaryFromLines("""
+blockimgdiff_versions=3,4
+use_dynamic_partitions=true
+dynamic_partition_list=foo
+super_partition_groups=group_foo
+super_group_foo_group_size={group_foo_size}
+super_group_foo_partition_list=foo
+""".format(group_foo_size=4 * GiB).split("\n"))
+    target_info = common.LoadDictionaryFromLines("""
+blockimgdiff_versions=3,4
+use_dynamic_partitions=true
+super_partition_groups=group_foo
+super_group_foo_group_size={group_foo_size}
+""".format(group_foo_size=4 * GiB).split("\n"))
+
+    common.OPTIONS.info_dict = target_info
+    common.OPTIONS.target_info_dict = target_info
+    common.OPTIONS.source_info_dict = source_info
+    common.OPTIONS.cache_size = 4 * 4096
+
+    block_diffs = [common.BlockDifference("foo", EmptyImage(),
+                                          src=DataImage("source", pad=True))]
+
+    dp_diff = common.DynamicPartitionsDifference(target_info, block_diffs,
+                                                 source_info_dict=source_info)
+    with zipfile.ZipFile(self.output_path, 'w') as output_zip:
+      dp_diff.WriteScript(self.script, output_zip, write_verify_script=True)
+
+    self.assertNotIn("block_image_update", str(self.script),
+        "Removed partition should not be patched.")
+
+    lines = self.get_op_list(self.output_path)
+    self.assertEqual(lines, ["remove foo"])