From 96527f375f734cf0e252ac65cdda0fb1807a3fd1 Mon Sep 17 00:00:00 2001
From: Yifan Hong <elsk@google.com>
Date: Mon, 28 Oct 2019 14:52:16 -0700
Subject: [PATCH] Fix Virtual A/B size checks

Test: build
Test: test_check_partition_sizes
Bug: 143111912
Change-Id: I4e056c25948e4169a0b5b098168141e27c31a0d4
---
 tools/releasetools/check_partition_sizes.py   | 14 ++++++++
 .../test_check_partition_sizes.py             | 34 +++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/tools/releasetools/check_partition_sizes.py b/tools/releasetools/check_partition_sizes.py
index 04d832ced..745c136e9 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 5482b1ceb..ed20873ac 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)