From 3612c88ed72cbe7a64696a94ca60a84b33fa13b1 Mon Sep 17 00:00:00 2001
From: Tao Bao <tbao@google.com>
Date: Mon, 14 Oct 2019 17:49:31 -0700
Subject: [PATCH] releasetools: Fix an issue in common.GetAvbPartitionArg.

It's a bug introduced in commit 1aeef725a7, which affects unittest only.

Bug: 130351427
Test: Use a lunch'd target. Run
      `atest --host releasetools_test releasetools_py3_test`.
Change-Id: I7ff01a6af47d002e1203bd376d477b60d769cbd1
---
 tools/releasetools/common.py      |  2 +-
 tools/releasetools/test_common.py | 33 +++++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index f846b18ff3..031db1d362 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -882,7 +882,7 @@ def GetAvbPartitionArg(partition, image, info_dict=None):
   # boot chain, to be verified with AVB_SLOT_VERIFY_FLAGS_NO_VBMETA_PARTITION.
   # See details at
   # https://android.googlesource.com/platform/external/avb/+/master/README.md#booting-into-recovery.
-  if OPTIONS.info_dict.get("ab_update") != "true" and partition == "recovery":
+  if info_dict.get("ab_update") != "true" and partition == "recovery":
     return []
 
   # Otherwise chain the partition into vbmeta.
diff --git a/tools/releasetools/test_common.py b/tools/releasetools/test_common.py
index 59b05e9d95..8a524194db 100644
--- a/tools/releasetools/test_common.py
+++ b/tools/releasetools/test_common.py
@@ -1373,6 +1373,39 @@ class CommonUtilsTest(test_utils.ReleaseToolsTestCase):
     self.assertEqual('5', chained_partition_args[1])
     self.assertTrue(os.path.exists(chained_partition_args[2]))
 
+  @test_utils.SkipIfExternalToolsUnavailable()
+  def test_AppendVBMetaArgsForPartition_recoveryAsChainedPartition_nonAb(self):
+    testdata_dir = test_utils.get_testdata_dir()
+    pubkey = os.path.join(testdata_dir, 'testkey.pubkey.pem')
+    info_dict = {
+        'avb_avbtool': 'avbtool',
+        'avb_recovery_key_path': pubkey,
+        'avb_recovery_rollback_index_location': 3,
+    }
+    cmd = common.GetAvbPartitionArg(
+        'recovery', '/path/to/recovery.img', info_dict)
+    self.assertFalse(cmd)
+
+  @test_utils.SkipIfExternalToolsUnavailable()
+  def test_AppendVBMetaArgsForPartition_recoveryAsChainedPartition_ab(self):
+    testdata_dir = test_utils.get_testdata_dir()
+    pubkey = os.path.join(testdata_dir, 'testkey.pubkey.pem')
+    info_dict = {
+        'ab_update': 'true',
+        'avb_avbtool': 'avbtool',
+        'avb_recovery_key_path': pubkey,
+        'avb_recovery_rollback_index_location': 3,
+    }
+    cmd = common.GetAvbPartitionArg(
+        'recovery', '/path/to/recovery.img', info_dict)
+    self.assertEqual(2, len(cmd))
+    self.assertEqual('--chain_partition', cmd[0])
+    chained_partition_args = cmd[1].split(':')
+    self.assertEqual(3, len(chained_partition_args))
+    self.assertEqual('recovery', chained_partition_args[0])
+    self.assertEqual('3', chained_partition_args[1])
+    self.assertTrue(os.path.exists(chained_partition_args[2]))
+
 
 class InstallRecoveryScriptFormatTest(test_utils.ReleaseToolsTestCase):
   """Checks the format of install-recovery.sh.