diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index 3e5dd52b0d..235883a0e0 100755
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -358,29 +358,39 @@ def AddUserdata(output_zip):
   img.Write()
 
 
-def AppendVBMetaArgsForPartition(cmd, partition, img_path, public_key_dir):
-  if not img_path:
-    return
+def AppendVBMetaArgsForPartition(cmd, partition, image):
+  """Appends the VBMeta arguments for partition.
 
+  It sets up the VBMeta argument by including the partition descriptor from the
+  given 'image', or by configuring the partition as a chained partition.
+
+  Args:
+    cmd: A list of command args that will be used to generate the vbmeta image.
+        The argument for the partition will be appended to the list.
+    partition: The name of the partition (e.g. "system").
+    image: The path to the partition image.
+  """
   # Check if chain partition is used.
   key_path = OPTIONS.info_dict.get("avb_" + partition + "_key_path")
   if key_path:
     # extract public key in AVB format to be included in vbmeta.img
     avbtool = os.getenv('AVBTOOL') or OPTIONS.info_dict["avb_avbtool"]
-    public_key_path = os.path.join(public_key_dir, "%s.avbpubkey" % partition)
-    p = common.Run([avbtool, "extract_public_key", "--key", key_path,
-                    "--output", public_key_path],
-                   stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    p.communicate()
-    assert p.returncode == 0, \
-        "avbtool extract_public_key fail for partition: %r" % partition
+    pubkey_path = common.MakeTempFile(prefix="avb-", suffix=".pubkey")
+    proc = common.Run(
+        [avbtool, "extract_public_key", "--key", key_path, "--output",
+         pubkey_path],
+        stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    stdoutdata, _ = proc.communicate()
+    assert proc.returncode == 0, \
+        "Failed to extract pubkey for {}:\n{}".format(
+            partition, stdoutdata)
 
     rollback_index_location = OPTIONS.info_dict[
         "avb_" + partition + "_rollback_index_location"]
     cmd.extend(["--chain_partition", "%s:%s:%s" % (
-        partition, rollback_index_location, public_key_path)])
+        partition, rollback_index_location, pubkey_path)])
   else:
-    cmd.extend(["--include_descriptors_from_image", img_path])
+    cmd.extend(["--include_descriptors_from_image", image])
 
 
 def AddVBMeta(output_zip, partitions):
@@ -389,8 +399,8 @@ def AddVBMeta(output_zip, partitions):
   Args:
     output_zip: The output zip file, which needs to be already open.
     partitions: A dict that's keyed by partition names with image paths as
-        values. Only valid partition names are accepted, which include 'boot',
-        'recovery', 'system', 'vendor', 'dtbo'.
+        values. Only valid partition names are accepted, as listed in
+        common.AVB_PARTITIONS.
   """
   img = OutputFile(output_zip, OPTIONS.input_tmp, "IMAGES", "vbmeta.img")
   if os.path.exists(img.input_name):
@@ -401,13 +411,12 @@ def AddVBMeta(output_zip, partitions):
   cmd = [avbtool, "make_vbmeta_image", "--output", img.name]
   common.AppendAVBSigningArgs(cmd, "vbmeta")
 
-  public_key_dir = common.MakeTempDir(prefix="avbpubkey-")
   for partition, path in partitions.items():
-    assert partition in common.AVB_PARTITIONS, 'Unknown partition: %s' % (
-        partition,)
-    assert os.path.exists(path), 'Failed to find %s for partition %s' % (
-        path, partition)
-    AppendVBMetaArgsForPartition(cmd, partition, path, public_key_dir)
+    assert partition in common.AVB_PARTITIONS, \
+        'Unknown partition: {}'.format(partition)
+    assert os.path.exists(path), \
+        'Failed to find {} for {}'.format(path, partition)
+    AppendVBMetaArgsForPartition(cmd, partition, path)
 
   args = OPTIONS.info_dict.get("avb_vbmeta_args")
   if args and args.strip():
diff --git a/tools/releasetools/test_add_img_to_target_files.py b/tools/releasetools/test_add_img_to_target_files.py
index a9c7f7caee..3f2e5ea451 100644
--- a/tools/releasetools/test_add_img_to_target_files.py
+++ b/tools/releasetools/test_add_img_to_target_files.py
@@ -22,7 +22,8 @@ import zipfile
 import common
 import test_utils
 from add_img_to_target_files import (
-    AddCareMapTxtForAbOta, AddPackRadioImages, CheckAbOtaImages, GetCareMap)
+    AddCareMapTxtForAbOta, AddPackRadioImages, AppendVBMetaArgsForPartition,
+    CheckAbOtaImages, GetCareMap)
 from rangelib import RangeSet
 
 
@@ -264,6 +265,31 @@ class AddImagesToTargetFilesTest(unittest.TestCase):
     # The existing entry should be scheduled to be replaced.
     self.assertIn('META/care_map.txt', OPTIONS.replace_updated_files_list)
 
+  def test_AppendVBMetaArgsForPartition(self):
+    OPTIONS.info_dict = {}
+    cmd = []
+    AppendVBMetaArgsForPartition(cmd, 'system', '/path/to/system.img')
+    self.assertEqual(
+        ['--include_descriptors_from_image', '/path/to/system.img'], cmd)
+
+  def test_AppendVBMetaArgsForPartition_vendorAsChainedPartition(self):
+    testdata_dir = test_utils.get_testdata_dir()
+    pubkey = os.path.join(testdata_dir, 'testkey.pubkey.pem')
+    OPTIONS.info_dict = {
+        'avb_avbtool': 'avbtool',
+        'avb_vendor_key_path': pubkey,
+        'avb_vendor_rollback_index_location': 5,
+    }
+    cmd = []
+    AppendVBMetaArgsForPartition(cmd, 'vendor', '/path/to/vendor.img')
+    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('vendor', chained_partition_args[0])
+    self.assertEqual('5', chained_partition_args[1])
+    self.assertTrue(os.path.exists(chained_partition_args[2]))
+
   def test_GetCareMap(self):
     sparse_image = test_utils.construct_sparse_image([
         (0xCAC1, 6),