Reland "Move partition size checks to python script"

This reverts commit 0141e45b9633fa21f1a2302d79db97d319ce7ec4.

Reason for revert: Fixed in follow-up CL
Test: forrest
Bug: 143734706

Change-Id: I007acf228d4fb4d6a16ae9089e3f04cf33a567bb
diff --git a/core/Makefile b/core/Makefile
index 6015796..96dc613 100644
--- a/core/Makefile
+++ b/core/Makefile
@@ -3616,36 +3616,6 @@
 
 ifeq (true,$(PRODUCT_BUILD_SUPER_PARTITION))
 
-# (1): list of items like "system", "vendor", "product", "system_ext"
-# return: map each item into a command ( wrapped in $$() ) that reads the size
-define read-size-of-partitions
-$(foreach image,$(call images-for-partitions,$(1)),$$($(SPARSE_IMG) --get_partition_size $(image)))
-endef
-
-# round result to BOARD_SUPER_PARTITION_ALIGNMENT
-#$(1): the calculated size
-ifeq (,$(BOARD_SUPER_PARTITION_ALIGNMENT))
-define round-partition-size
-$(1)
-endef
-else
-define round-partition-size
-$$((($(1)+$(BOARD_SUPER_PARTITION_ALIGNMENT)-1)/$(BOARD_SUPER_PARTITION_ALIGNMENT)*$(BOARD_SUPER_PARTITION_ALIGNMENT)))
-endef
-endif
-
-define super-slot-suffix
-$(if $(filter true,$(AB_OTA_UPDATER)),$(if $(filter true,$(PRODUCT_RETROFIT_DYNAMIC_PARTITIONS)),,_a))
-endef
-
-ifndef BOARD_SUPER_PARTITION_WARN_LIMIT
-BOARD_SUPER_PARTITION_WARN_LIMIT := $$(($(BOARD_SUPER_PARTITION_SIZE) * 95 / 100))
-endif
-
-ifndef BOARD_SUPER_PARTITION_ERROR_LIMIT
-BOARD_SUPER_PARTITION_ERROR_LIMIT := $(BOARD_SUPER_PARTITION_SIZE)
-endif
-
 droid_targets: check-all-partition-sizes
 
 .PHONY: check-all-partition-sizes check-all-partition-sizes-nodeps
@@ -3654,103 +3624,28 @@
 
 check-all-partition-sizes: $(check_all_partition_sizes_file)
 
-# Add image dependencies so that generated_*_image_info.txt are written before checking.
 $(check_all_partition_sizes_file): \
-    $(SPARSE_IMG) \
+    $(CHECK_PARTITION_SIZES) \
     $(call images-for-partitions,$(BOARD_SUPER_PARTITION_PARTITION_LIST))
 
-ifeq ($(PRODUCT_RETROFIT_DYNAMIC_PARTITIONS),true)
-# Check sum(super partition block devices) == super partition
-# Non-retrofit devices already defines BOARD_SUPER_PARTITION_SUPER_DEVICE_SIZE = BOARD_SUPER_PARTITION_SIZE
-define check-super-partition-size
-  size_list="$(foreach device,$(call to-upper,$(BOARD_SUPER_PARTITION_BLOCK_DEVICES)),$(BOARD_SUPER_PARTITION_$(device)_DEVICE_SIZE))"; \
-  sum_sizes_expr=$$(sed -e 's/ /+/g' <<< "$${size_list}"); \
-  max_size_expr="$(BOARD_SUPER_PARTITION_SIZE)"; \
-  if [ $$(( $${sum_sizes_expr} )) -ne $$(( $${max_size_expr} )) ]; then \
-    echo "The sum of super partition block device sizes is not equal to BOARD_SUPER_PARTITION_SIZE:"; \
-    echo $${sum_sizes_expr} '!=' $${max_size_expr}; \
-    exit 1; \
-  else \
-    echo "The sum of super partition block device sizes is equal to BOARD_SUPER_PARTITION_SIZE:"; \
-    echo $${sum_sizes_expr} '==' $${max_size_expr}; \
-  fi
-endef
-endif
-
-# $(1): human-readable max size string
-# $(2): max size expression
-# $(3): list of partition names
-# $(4): human-readable warn size string
-# $(5): warn size expression
-# $(6): human readable error size string
-# $(7): error size expression
-define check-sum-of-partition-sizes
-  partition_size_list="$$(for i in $(call read-size-of-partitions,$(3)); do \
-    echo $(call round-partition-size,$${i}); \
-  done)"; \
-  sum_sizes_expr=$$(tr '\n' '+' <<< "$${partition_size_list}" | sed 's/+$$//'); \
-  if [ $$(( $${sum_sizes_expr} )) -gt $$(( $(2) )) ]; then \
-    echo "The sum of sizes of [$(strip $(3))] is larger than $(strip $(1)):"; \
-    echo $${sum_sizes_expr} '==' $$(( $${sum_sizes_expr} )) '>' "$(2)" '==' $$(( $(2) )); \
-    exit 1; \
-  else \
-    if [[ ! -z "$(7)" ]] && [ $$(( $${sum_sizes_expr} )) -gt $$(( $(7) )) ]; then \
-        echo "!!!! ERROR !!!! The sum of sizes of [$(strip $(3))] is larger than $(strip $(6)):"; \
-        echo $${sum_sizes_expr} '==' $$(( $${sum_sizes_expr} )) '>' "$(7)" '==' $$(( $(7) )); \
-        echo "Super partition is" $$(( $$(( $$(( $${sum_sizes_expr} )) * 100)) / $$(( $(2) )) )) "percent occupied!"; \
-        exit 1; \
-    fi; \
-    if [[ ! -z "$(5)" ]] && [ $$(( $${sum_sizes_expr} )) -gt $$(( $(5) )) ]; then \
-        echo "!!!! WARNING !!!! The sum of sizes of [$(strip $(3))] is larger than $(strip $(4)):"; \
-        echo $${sum_sizes_expr} '==' $$(( $${sum_sizes_expr} )) '>' "$(5)" '==' $$(( $(5) )); \
-        echo "Super partition is" $$(( $$(( $$(( $${sum_sizes_expr} )) * 100)) / $$(( $(2) )) )) "percent occupied!"; \
-    fi; \
-    echo "The sum of sizes of [$(strip $(3))] is within $(strip $(1)):"; \
-    echo $${sum_sizes_expr} '==' $$(( $${sum_sizes_expr} )) '<=' "$(2)" '==' $$(( $(2) )); \
-  fi;
-endef
-
+# $(1): misc_info.txt
 define check-all-partition-sizes-target
-  # Check sum(all partitions) <= super partition (/ 2 for A/B devices launched with dynamic partitions)
-  $(if $(BOARD_SUPER_PARTITION_SIZE),$(if $(BOARD_SUPER_PARTITION_PARTITION_LIST), \
-    $(call check-sum-of-partition-sizes,BOARD_SUPER_PARTITION_SIZE$(if $(call super-slot-suffix), / 2), \
-      $(BOARD_SUPER_PARTITION_SIZE)$(if $(call super-slot-suffix), / 2),$(BOARD_SUPER_PARTITION_PARTITION_LIST), \
-      BOARD_SUPER_PARTITION_WARN_LIMIT$(if $(call super-slot-suffix), / 2), \
-      $(BOARD_SUPER_PARTITION_WARN_LIMIT)$(if $(call super-slot-suffix), / 2), \
-      BOARD_SUPER_PARTITION_ERROR_LIMIT$(if $(call super-slot-suffix), / 2), \
-      $(BOARD_SUPER_PARTITION_ERROR_LIMIT)$(if $(call super-slot-suffix), / 2)) \
-  ))
-
-  # For each group, check sum(partitions in group) <= group size
-  $(foreach group,$(call to-upper,$(BOARD_SUPER_PARTITION_GROUPS)), \
-    $(if $(BOARD_$(group)_SIZE),$(if $(BOARD_$(group)_PARTITION_LIST), \
-      $(call check-sum-of-partition-sizes,BOARD_$(group)_SIZE,$(BOARD_$(group)_SIZE),$(BOARD_$(group)_PARTITION_LIST)))))
-
-  # Check sum(all group sizes) <= super partition (/ 2 for A/B devices launched with dynamic partitions)
-  if [[ ! -z $(BOARD_SUPER_PARTITION_SIZE) ]]; then \
-    group_size_list="$(foreach group,$(call to-upper,$(BOARD_SUPER_PARTITION_GROUPS)),$(BOARD_$(group)_SIZE))"; \
-    sum_sizes_expr=$$(sed -e 's/ /+/g' <<< "$${group_size_list}"); \
-    max_size_tail=$(if $(call super-slot-suffix)," / 2"); \
-    max_size_expr="$(BOARD_SUPER_PARTITION_SIZE)$${max_size_tail}"; \
-    if [ $$(( $${sum_sizes_expr} )) -gt $$(( $${max_size_expr} )) ]; then \
-      echo "The sum of sizes of [$(strip $(BOARD_SUPER_PARTITION_GROUPS))] is larger than BOARD_SUPER_PARTITION_SIZE$${max_size_tail}:"; \
-      echo $${sum_sizes_expr} '==' $$(( $${sum_sizes_expr} )) '>' $${max_size_expr} '==' $$(( $${max_size_expr} )); \
-      exit 1; \
-    else \
-      echo "The sum of sizes of [$(strip $(BOARD_SUPER_PARTITION_GROUPS))] is within BOARD_SUPER_PARTITION_SIZE$${max_size_tail}:"; \
-      echo $${sum_sizes_expr} '==' $$(( $${sum_sizes_expr} )) '<=' $${max_size_expr} '==' $$(( $${max_size_expr} )); \
-    fi \
-  fi
+  mkdir -p $(dir $(1))
+  rm -f $(1)
+  $(call dump-super-image-info, $(1))
+  $(foreach partition,$(BOARD_SUPER_PARTITION_PARTITION_LIST), \
+    echo "$(partition)_image="$(call images-for-partitions,$(partition)) >> $(1);)
+  $(CHECK_PARTITION_SIZES) -v $(1)
 endef
 
 $(check_all_partition_sizes_file):
-	$(call check-all-partition-sizes-target)
-	$(call check-super-partition-size)
+	$(call check-all-partition-sizes-target, \
+	  $(call intermediates-dir-for,PACKAGING,check-all-partition-sizes)/misc_info.txt)
 	touch $@
 
 check-all-partition-sizes-nodeps:
-	$(call check-all-partition-sizes-target)
-	$(call check-super-partition-size)
+	$(call check-all-partition-sizes-target, \
+	  $(call intermediates-dir-for,PACKAGING,check-all-partition-sizes-nodeps)/misc_info.txt)
 
 endif # PRODUCT_BUILD_SUPER_PARTITION
 
@@ -4189,6 +4084,14 @@
     echo "build_non_sparse_super_partition=true" >> $(1))
   $(if $(filter true,$(BOARD_SUPER_IMAGE_IN_UPDATE_PACKAGE)), \
     echo "super_image_in_update_package=true" >> $(1))
+  $(if $(BOARD_SUPER_PARTITION_SIZE), \
+    echo "super_partition_size=$(BOARD_SUPER_PARTITION_SIZE)" >> $(1))
+  $(if $(BOARD_SUPER_PARTITION_ALIGNMENT), \
+    echo "super_partition_alignment=$(BOARD_SUPER_PARTITION_ALIGNMENT)" >> $(1))
+  $(if $(BOARD_SUPER_PARTITION_WARN_LIMIT), \
+    echo "super_partition_warn_limit=$(BOARD_SUPER_PARTITION_WARN_LIMIT)" >> $(1))
+  $(if $(BOARD_SUPER_PARTITION_ERROR_LIMIT), \
+    echo "super_partition_error_limit=$(BOARD_SUPER_PARTITION_ERROR_LIMIT)" >> $(1))
 endef
 
 # By conditionally including the dependency of the target files package on the
@@ -4765,6 +4668,10 @@
   $(call dump-dynamic-partitions-info,$(1))
   $(if $(filter true,$(AB_OTA_UPDATER)), \
     echo "ab_update=true" >> $(1))
+  $(if $(filter true,$(PRODUCT_VIRTUAL_AB_OTA)), \
+    echo "virtual_ab=true" >> $(1))
+  $(if $(filter true,$(PRODUCT_VIRTUAL_AB_OTA_RETROFIT)), \
+    echo "virtual_ab_retrofit=true" >> $(1))
 endef
 
 endif # PRODUCT_USE_DYNAMIC_PARTITIONS
diff --git a/core/config.mk b/core/config.mk
index 477d764..a1bbe18 100644
--- a/core/config.mk
+++ b/core/config.mk
@@ -602,6 +602,7 @@
 MAKE_RECOVERY_PATCH := $(HOST_OUT_EXECUTABLES)/make_recovery_patch$(HOST_EXECUTABLE_SUFFIX)
 OTA_FROM_TARGET_FILES := $(HOST_OUT_EXECUTABLES)/ota_from_target_files$(HOST_EXECUTABLE_SUFFIX)
 SPARSE_IMG := $(HOST_OUT_EXECUTABLES)/sparse_img$(HOST_EXECUTABLE_SUFFIX)
+CHECK_PARTITION_SIZES := $(HOST_OUT_EXECUTABLES)/check_partition_sizes$(HOST_EXECUTABLE_SUFFIX)
 
 PROGUARD_HOME := external/proguard
 PROGUARD := $(PROGUARD_HOME)/bin/proguard.sh
diff --git a/tools/releasetools/Android.bp b/tools/releasetools/Android.bp
index 6cde77e..90a6485 100644
--- a/tools/releasetools/Android.bp
+++ b/tools/releasetools/Android.bp
@@ -264,6 +264,19 @@
 }
 
 python_binary_host {
+    name: "check_partition_sizes",
+    srcs: [
+        "check_partition_sizes.py",
+    ],
+    libs: [
+        "releasetools_common",
+    ],
+    defaults: [
+        "releasetools_binary_defaults",
+    ],
+}
+
+python_binary_host {
     name: "check_ota_package_signature",
     defaults: ["releasetools_binary_defaults"],
     srcs: [
@@ -419,6 +432,7 @@
     name: "releasetools_test_defaults",
     srcs: [
         "check_ota_package_signature.py",
+        "check_partition_sizes.py",
         "check_target_files_signatures.py",
         "make_recovery_patch.py",
         "merge_target_files.py",
diff --git a/tools/releasetools/check_partition_sizes.py b/tools/releasetools/check_partition_sizes.py
new file mode 100644
index 0000000..04d832c
--- /dev/null
+++ b/tools/releasetools/check_partition_sizes.py
@@ -0,0 +1,262 @@
+#!/usr/bin/env python
+#
+# Copyright (C) 2019 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.
+
+"""
+Check dynamic partition sizes.
+
+usage: check_partition_sizes [info.txt]
+
+Check dump-super-partitions-info procedure for expected keys in info.txt. In
+addition, *_image (e.g. system_image, vendor_image, etc.) must be defined for
+each partition in dynamic_partition_list.
+
+Exit code is 0 if successful and non-zero if any failures.
+"""
+
+from __future__ import print_function
+
+import logging
+import sys
+
+import common
+import sparse_img
+
+if sys.hexversion < 0x02070000:
+  print("Python 2.7 or newer is required.", file=sys.stderr)
+  sys.exit(1)
+
+logger = logging.getLogger(__name__)
+
+class Expression(object):
+  def __init__(self, desc, expr, value=None):
+    # Human-readable description
+    self.desc = str(desc)
+    # Numeric expression
+    self.expr = str(expr)
+    # Value of expression
+    self.value = int(expr) if value is None else value
+
+  def CheckLe(self, other, level=logging.ERROR):
+    format_args = (self.desc, other.desc, self.expr, self.value,
+                   other.expr, other.value)
+    if self.value <= other.value:
+      logger.info("%s is less than or equal to %s:\n%s == %d <= %s == %d",
+                  *format_args)
+    else:
+      msg = "{} is greater than {}:\n{} == {} > {} == {}".format(*format_args)
+      if level == logging.ERROR:
+        raise RuntimeError(msg)
+      else:
+        logger.log(level, msg)
+
+  def CheckEq(self, other):
+    format_args = (self.desc, other.desc, self.expr, self.value,
+                   other.expr, other.value)
+    if self.value == other.value:
+      logger.info("%s equals %s:\n%s == %d == %s == %d", *format_args)
+    else:
+      raise RuntimeError("{} does not equal {}:\n{} == {} != {} == {}".format(
+          *format_args))
+
+
+# A/B feature flags
+class DeviceType(object):
+  NONE = 0
+  AB = 1
+
+  @staticmethod
+  def Get(info_dict):
+    if info_dict.get("ab_update") != "true":
+      return DeviceType.NONE
+    return DeviceType.AB
+
+
+# Dynamic partition feature flags
+class Dap(object):
+  NONE = 0
+  RDAP = 1
+  DAP = 2
+
+  @staticmethod
+  def Get(info_dict):
+    if info_dict.get("use_dynamic_partitions") != "true":
+      return Dap.NONE
+    if info_dict.get("dynamic_partition_retrofit") == "true":
+      return Dap.RDAP
+    return Dap.DAP
+
+
+class DynamicPartitionSizeChecker(object):
+  def __init__(self, info_dict):
+    if "super_partition_size" in info_dict:
+      if "super_partition_warn_limit" not in info_dict:
+        info_dict["super_partition_warn_limit"] = \
+            int(info_dict["super_partition_size"]) * 95 // 100
+      if "super_partition_error_limit" not in info_dict:
+        info_dict["super_partition_error_limit"] = \
+            int(info_dict["super_partition_size"])
+    self.info_dict = info_dict
+
+
+  def _ReadSizeOfPartition(self, name):
+    # Tests uses *_image_size instead (to avoid creating empty sparse images
+    # on disk)
+    if name + "_image_size" in self.info_dict:
+      return int(self.info_dict[name + "_image_size"])
+    return sparse_img.GetImagePartitionSize(self.info_dict[name + "_image"])
+
+
+  # Round result to BOARD_SUPER_PARTITION_ALIGNMENT
+  def _RoundPartitionSize(self, size):
+    alignment = self.info_dict.get("super_partition_alignment")
+    if alignment is None:
+      return size
+    return (size + alignment - 1) // alignment * alignment
+
+
+  def _CheckSuperPartitionSize(self):
+    info_dict = self.info_dict
+    super_block_devices = \
+        info_dict.get("super_block_devices", "").strip().split()
+    size_list = [int(info_dict.get("super_{}_device_size".format(b), "0"))
+                 for b in super_block_devices]
+    sum_size = Expression("sum of super partition block device sizes",
+                          "+".join(str(size) for size in size_list),
+                          sum(size_list))
+    super_partition_size = Expression("BOARD_SUPER_PARTITION_SIZE",
+                                      info_dict["super_partition_size"])
+    sum_size.CheckEq(super_partition_size)
+
+  def _CheckSumOfPartitionSizes(self, max_size, partition_names,
+                                warn_size=None, error_size=None):
+    partition_size_list = [self._RoundPartitionSize(
+        self._ReadSizeOfPartition(p)) for p in partition_names]
+    sum_size = Expression("sum of sizes of {}".format(partition_names),
+                          "+".join(str(size) for size in partition_size_list),
+                          sum(partition_size_list))
+    sum_size.CheckLe(max_size)
+    if error_size:
+      sum_size.CheckLe(error_size)
+    if warn_size:
+      sum_size.CheckLe(warn_size, level=logging.WARNING)
+
+  def _NumDeviceTypesInSuper(self):
+    slot = DeviceType.Get(self.info_dict)
+    dap = Dap.Get(self.info_dict)
+
+    if dap == Dap.NONE:
+      raise RuntimeError("check_partition_sizes should only be executed on "
+                         "builds with dynamic partitions enabled")
+
+    # Retrofit dynamic partitions: 1 slot per "super", 2 "super"s on the device
+    if dap == Dap.RDAP:
+      if slot != DeviceType.AB:
+        raise RuntimeError("Device with retrofit dynamic partitions must use "
+                           "regular (non-Virtual) A/B")
+      return 1
+
+    # Launch DAP: 1 super on the device
+    assert dap == Dap.DAP
+
+    # DAP + A/B: 2 slots in super
+    if slot == DeviceType.AB:
+      return 2
+
+    # DAP + non-A/B: 1 slot in super
+    assert slot == DeviceType.NONE
+    return 1
+
+  def _CheckAllPartitionSizes(self):
+    info_dict = self.info_dict
+    num_slots = self._NumDeviceTypesInSuper()
+    size_limit_suffix = (" / %d" % num_slots) if num_slots > 1 else ""
+
+    # Check sum(all partitions) <= super partition (/ 2 for A/B devices launched
+    # with dynamic partitions)
+    if "super_partition_size" in info_dict and \
+        "dynamic_partition_list" in info_dict:
+      max_size = Expression(
+          "BOARD_SUPER_PARTITION_SIZE{}".format(size_limit_suffix),
+          int(info_dict["super_partition_size"]) // num_slots)
+      warn_limit = Expression(
+          "BOARD_SUPER_PARTITION_WARN_LIMIT{}".format(size_limit_suffix),
+          int(info_dict["super_partition_warn_limit"]) // num_slots)
+      error_limit = Expression(
+          "BOARD_SUPER_PARTITION_ERROR_LIMIT{}".format(size_limit_suffix),
+          int(info_dict["super_partition_error_limit"]) // num_slots)
+      self._CheckSumOfPartitionSizes(
+          max_size, info_dict["dynamic_partition_list"].strip().split(),
+          warn_limit, error_limit)
+
+    groups = info_dict.get("super_partition_groups", "").strip().split()
+
+    # For each group, check sum(partitions in group) <= group size
+    for group in groups:
+      if "super_{}_group_size".format(group) in info_dict and \
+          "super_{}_partition_list".format(group) in info_dict:
+        group_size = Expression(
+            "BOARD_{}_SIZE".format(group),
+            int(info_dict["super_{}_group_size".format(group)]))
+        self._CheckSumOfPartitionSizes(
+            group_size,
+            info_dict["super_{}_partition_list".format(group)].strip().split())
+
+    # Check sum(all group sizes) <= super partition (/ 2 for A/B devices
+    # launched with dynamic partitions)
+    if "super_partition_size" in info_dict:
+      group_size_list = [int(info_dict.get(
+          "super_{}_group_size".format(group), 0)) for group in groups]
+      sum_size = Expression("sum of sizes of {}".format(groups),
+                            "+".join(str(size) for size in group_size_list),
+                            sum(group_size_list))
+      max_size = Expression(
+          "BOARD_SUPER_PARTITION_SIZE{}".format(size_limit_suffix),
+          int(info_dict["super_partition_size"]) // num_slots)
+      sum_size.CheckLe(max_size)
+
+  def Run(self):
+    self._CheckAllPartitionSizes()
+    if self.info_dict.get("dynamic_partition_retrofit") == "true":
+      self._CheckSuperPartitionSize()
+
+
+def CheckPartitionSizes(inp):
+  if isinstance(inp, str):
+    info_dict = common.LoadDictionaryFromFile(inp)
+    return DynamicPartitionSizeChecker(info_dict).Run()
+  if isinstance(inp, dict):
+    return DynamicPartitionSizeChecker(inp).Run()
+  raise ValueError("{} is not a dictionary or a valid path".format(inp))
+
+
+def main(argv):
+  args = common.ParseOptions(argv, __doc__)
+  if len(args) != 1:
+    common.Usage(__doc__)
+    sys.exit(1)
+  common.InitLogging()
+  CheckPartitionSizes(args[0])
+
+
+if __name__ == "__main__":
+  try:
+    common.CloseInheritedPipes()
+    main(sys.argv[1:])
+  except common.ExternalError:
+    logger.exception("\n   ERROR:\n")
+    sys.exit(1)
+  finally:
+    common.Cleanup()
diff --git a/tools/releasetools/test_check_partition_sizes.py b/tools/releasetools/test_check_partition_sizes.py
new file mode 100644
index 0000000..5482b1c
--- /dev/null
+++ b/tools/releasetools/test_check_partition_sizes.py
@@ -0,0 +1,94 @@
+#
+# Copyright (C) 2019 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.
+#
+
+import common
+import test_utils
+from check_partition_sizes import CheckPartitionSizes
+
+class CheckPartitionSizesTest(test_utils.ReleaseToolsTestCase):
+  def setUp(self):
+    self.info_dict = common.LoadDictionaryFromLines("""
+        use_dynamic_partitions=true
+        ab_update=true
+        super_block_devices=super
+        dynamic_partition_list=system vendor product
+        super_partition_groups=group
+        super_group_partition_list=system vendor product
+        super_partition_size=200
+        super_super_device_size=200
+        super_group_group_size=100
+        system_image_size=50
+        vendor_image_size=20
+        product_image_size=20
+        """.split("\n"))
+
+  def test_ab(self):
+    CheckPartitionSizes(self.info_dict)
+
+  def test_non_ab(self):
+    self.info_dict.update(common.LoadDictionaryFromLines("""
+        ab_update=false
+        super_partition_size=100
+        super_super_device_size=100
+        """.split("\n")))
+    CheckPartitionSizes(self.info_dict)
+
+  def test_non_dap(self):
+    self.info_dict.update(common.LoadDictionaryFromLines("""
+        use_dynamic_partitions=false
+        """.split("\n")))
+    with self.assertRaises(RuntimeError):
+      CheckPartitionSizes(self.info_dict)
+
+  def test_retrofit_dap(self):
+    self.info_dict.update(common.LoadDictionaryFromLines("""
+        dynamic_partition_retrofit=true
+        super_block_devices=system vendor
+        super_system_device_size=75
+        super_vendor_device_size=25
+        super_partition_size=100
+        """.split("\n")))
+    CheckPartitionSizes(self.info_dict)
+
+  def test_ab_partition_too_big(self):
+    self.info_dict.update(common.LoadDictionaryFromLines("""
+        system_image_size=100
+        """.split("\n")))
+    with self.assertRaises(RuntimeError):
+      CheckPartitionSizes(self.info_dict)
+
+  def test_ab_group_too_big(self):
+    self.info_dict.update(common.LoadDictionaryFromLines("""
+        super_group_group_size=110
+        """.split("\n")))
+    with self.assertRaises(RuntimeError):
+      CheckPartitionSizes(self.info_dict)
+
+  def test_no_image(self):
+    del self.info_dict["system_image_size"]
+    with self.assertRaises(KeyError):
+      CheckPartitionSizes(self.info_dict)
+
+  def test_block_devices_not_match(self):
+    self.info_dict.update(common.LoadDictionaryFromLines("""
+        dynamic_partition_retrofit=true
+        super_block_devices=system vendor
+        super_system_device_size=80
+        super_vendor_device_size=25
+        super_partition_size=100
+        """.split("\n")))
+    with self.assertRaises(RuntimeError):
+      CheckPartitionSizes(self.info_dict)