Fix Virtual A/B size checks

Test: build
Test: test_check_partition_sizes
Bug: 143111912
Change-Id: I4e056c25948e4169a0b5b098168141e27c31a0d4
diff --git a/tools/releasetools/check_partition_sizes.py b/tools/releasetools/check_partition_sizes.py
index 04d832c..745c136 100644
--- a/tools/releasetools/check_partition_sizes.py
+++ b/tools/releasetools/check_partition_sizes.py
@@ -76,11 +76,17 @@
 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 @@
     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