Merge "Fix Virtual A/B size checks"

2 files changed, 48 insertions, 0 deletions

diff --git a/tools/releasetools/check_partition_sizes.py b/tools/releasetools/check_partition_sizes.py
index 04d832ced8..745c136e90 100644
--- a/tools/releasetools/check_partition_sizes.py
+++ b/tools/releasetools/check_partition_sizes.py
@@ -76,11 +76,17 @@ class Expression(object):
 class DeviceType(object):
   NONE = 0
   AB = 1
+  RVAB = 2 # retrofit Virtual-A/B
+  VAB = 3
 
   @staticmethod
   def Get(info_dict):
     if info_dict.get("ab_update") != "true":
       return DeviceType.NONE
+    if info_dict.get("virtual_ab_retrofit") == "true":
+      return DeviceType.RVAB
+    if info_dict.get("virtual_ab") == "true":
+      return DeviceType.VAB
     return DeviceType.AB
 
 
@@ -175,6 +181,14 @@ class DynamicPartitionSizeChecker(object):
     if slot == DeviceType.AB:
       return 2
 
+    # DAP + retrofit Virtual A/B: same as A/B
+    if slot == DeviceType.RVAB:
+      return 2
+
+    # DAP + Launch Virtual A/B: 1 *real* slot in super (2 virtual slots)
+    if slot == DeviceType.VAB:
+      return 1
+
     # DAP + non-A/B: 1 slot in super
     assert slot == DeviceType.NONE
     return 1
diff --git a/tools/releasetools/test_check_partition_sizes.py b/tools/releasetools/test_check_partition_sizes.py
index 5482b1cebb..ed20873acb 100644
--- a/tools/releasetools/test_check_partition_sizes.py
+++ b/tools/releasetools/test_check_partition_sizes.py
@@ -92,3 +92,37 @@ class CheckPartitionSizesTest(test_utils.ReleaseToolsTestCase):
         """.split("\n")))
     with self.assertRaises(RuntimeError):
       CheckPartitionSizes(self.info_dict)
+
+  def test_retrofit_vab(self):
+    self.info_dict.update(common.LoadDictionaryFromLines("""
+        virtual_ab=true
+        virtual_ab_retrofit=true
+        """.split("\n")))
+    CheckPartitionSizes(self.info_dict)
+
+  def test_retrofit_vab_too_big(self):
+    self.info_dict.update(common.LoadDictionaryFromLines("""
+        virtual_ab=true
+        virtual_ab_retrofit=true
+        system_image_size=100
+        """.split("\n")))
+    with self.assertRaises(RuntimeError):
+      CheckPartitionSizes(self.info_dict)
+
+  def test_vab(self):
+    self.info_dict.update(common.LoadDictionaryFromLines("""
+        virtual_ab=true
+        super_partition_size=100
+        super_super_device_size=100
+        """.split("\n")))
+    CheckPartitionSizes(self.info_dict)
+
+  def test_vab_too_big(self):
+    self.info_dict.update(common.LoadDictionaryFromLines("""
+        virtual_ab=true
+        super_partition_size=100
+        super_super_device_size=100
+        system_image_size=100
+        """.split("\n")))
+    with self.assertRaises(RuntimeError):
+      CheckPartitionSizes(self.info_dict)