Add local-vendor-boot-image argument to set vendor boot image

To use debug vendor_boot image from local with remote image / remote
instance, it needs an option to specify local vendor boot image path.

Add an option local-vendor-boot-image for specifying vendor boot image.

Test: acloud-dev create --build-target target1 --branch branch1 \
        --build-id build1 --remote-fetch --host host1 \
	--host-user user1 -vv \
	--local-vendor-boot-image=vendor_boot-debug.img
Test: atest --host acloud_test
Bug: 335821657
Change-Id: Idbbc24b9e8edaf80bafc79a473fa58387ca561be

12 files changed, 148 insertions, 21 deletions

diff --git a/create/avd_spec.py b/create/avd_spec.py
index 8b2f984a..1879aaf0 100644
--- a/create/avd_spec.py
+++ b/create/avd_spec.py
@@ -124,6 +124,7 @@ class AVDSpec():
         self._local_system_image = None
         self._local_system_dlkm_image = None
         self._local_vendor_image = None
+        self._local_vendor_boot_image = None
         self._local_tool_dirs = None
         self._image_download_dir = None
         self._num_of_instances = None
@@ -269,6 +270,10 @@ class AVDSpec():
             self._local_vendor_image = self._GetLocalImagePath(
                 args.local_vendor_image)
 
+        if args.local_vendor_boot_image is not None:
+            self._local_vendor_boot_image = self._GetLocalImagePath(
+                args.local_vendor_boot_image)
+
         self.image_download_dir = (
             args.image_download_dir if args.image_download_dir
             else tempfile.gettempdir())
@@ -857,6 +862,11 @@ class AVDSpec():
         return self._local_vendor_image
 
     @property
+    def local_vendor_boot_image(self):
+        """Return local vendor boot image path."""
+        return self._local_vendor_boot_image
+
+    @property
     def local_tool_dirs(self):
         """Return a list of local tool directories."""
         return self._local_tool_dirs
diff --git a/create/avd_spec_test.py b/create/avd_spec_test.py
index 7fe45cca..173322bc 100644
--- a/create/avd_spec_test.py
+++ b/create/avd_spec_test.py
@@ -43,6 +43,7 @@ class AvdSpecTest(driver_test_lib.BaseDriverTest):
         self.args.local_image = None
         self.args.local_kernel_image = None
         self.args.local_system_image = None
+        self.args.local_vendor_boot_image = None
         self.args.config_file = ""
         self.args.build_target = "fake_build_target"
         self.args.adb_port = None
@@ -132,10 +133,12 @@ class AvdSpecTest(driver_test_lib.BaseDriverTest):
         self.args.local_kernel_image = expected_image_file
         self.args.local_system_image = expected_image_file
         self.args.local_system_dlkm_image = expected_image_file
+        self.args.local_vendor_boot_image = expected_image_file
         self.AvdSpec._ProcessImageArgs(self.args)
         self.assertEqual(self.AvdSpec.local_kernel_image, expected_image_file)
         self.assertEqual(self.AvdSpec.local_system_image, expected_image_file)
         self.assertEqual(self.AvdSpec.local_system_dlkm_image, expected_image_file)
+        self.assertEqual(self.AvdSpec.local_vendor_boot_image, expected_image_file)
 
         # Specified --local-*-image without args.
         self.args.local_kernel_image = constants.FIND_IN_BUILD_ENV
diff --git a/create/create_args.py b/create/create_args.py
index ce2e942e..3bd3dba2 100644
--- a/create/create_args.py
+++ b/create/create_args.py
@@ -627,6 +627,18 @@ def GetCreateArgParser(subparser):
         "$ANDROID_PRODUCT_OUT if no argument is provided. e.g., "
         "--local-vendor-image, or --local-vendor-image /path/to/dir")
     create_parser.add_argument(
+        "--local-vendor_boot-image", "--local-vendor-boot-image",
+        const=constants.FIND_IN_BUILD_ENV,
+        type=str,
+        dest="local_vendor_boot_image",
+        nargs="?",
+        required=False,
+        help="'cuttlefish only' Use the locally built vendor boot image for "
+        "the AVD. Look for the vendor_boot.img in $ANDROID_PRODUCT_OUT "
+        "if no argument is provided. e.g., --local-vendor-boot-image, or "
+        "--local-vendor-boot-image /path/to/dir, or "
+        "--local-vendor-boot-image /path/to/img")
+    create_parser.add_argument(
         "--local-tool",
         type=str,
         dest="local_tool",
diff --git a/create/create_args_test.py b/create/create_args_test.py
index 1364a448..279844b2 100644
--- a/create/create_args_test.py
+++ b/create/create_args_test.py
@@ -41,6 +41,7 @@ def _CreateArgs():
         local_kernel_image=None,
         local_system_image=None,
         local_instance_dir=None,
+        local_vendor_boot_image=None,
         kernel_branch=None,
         kernel_build_id=None,
         kernel_build_target="kernel",
diff --git a/create/create_common.py b/create/create_common.py
index f4dd5e0e..8897e6e1 100644
--- a/create/create_common.py
+++ b/create/create_common.py
@@ -40,6 +40,7 @@ _TARGET_FILES_IMAGES_DIR_NAME = "IMAGES"
 _SYSTEM_IMAGE_NAME = "system.img"
 _SYSTEM_EXT_IMAGE_NAME = "system_ext.img"
 _PRODUCT_IMAGE_NAME = "product.img"
+_VENDOR_BOOT_IMAGE_NAME_PATTERN = r"vendor_boot\.img"
 
 _ANDROID_BOOT_IMAGE_MAGIC = b"ANDROID!"
 
@@ -220,6 +221,11 @@ def FindBootImage(path, raise_error=True):
     return boot_image_path
 
 
+def FindVendorBootImage(path, raise_error=True):
+    """Find a vendor boot image file in the given path."""
+    return FindLocalImage(path, _VENDOR_BOOT_IMAGE_NAME_PATTERN, raise_error)
+
+
 def FindSystemImages(path):
     """Find system, system_ext, and product image files in a given path.
 
diff --git a/create/create_common_test.py b/create/create_common_test.py
index 84581857..52d361a7 100644
--- a/create/create_common_test.py
+++ b/create/create_common_test.py
@@ -218,6 +218,21 @@ class CreateCommonTest(driver_test_lib.BaseDriverTest):
                 (system_image_path, system_ext_image_path, product_image_path),
                 create_common.FindSystemImages(image_dir))
 
+
+    def testFindVendorBootImage(self):
+        """Test FindVendorBootImage."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            with self.assertRaises(errors.GetLocalImageError):
+                create_common.FindVendorBootImage(temp_dir)
+
+            boot_image_path = os.path.join(temp_dir, "vendor_boot.img")
+            self.CreateFile(boot_image_path)
+            self.assertEqual(boot_image_path,
+                             create_common.FindVendorBootImage(boot_image_path))
+            self.assertEqual(boot_image_path,
+                             create_common.FindVendorBootImage(temp_dir))
+
+
     @mock.patch.object(utils, "Decompress")
     def testDownloadRemoteArtifact(self, mock_decompress):
         """Test Download cuttlefish package."""
diff --git a/create/local_image_local_instance.py b/create/local_image_local_instance.py
index de6819f4..d00cf84d 100644
--- a/create/local_image_local_instance.py
+++ b/create/local_image_local_instance.py
@@ -522,6 +522,10 @@ class LocalImageLocalInstance(base_avd_create.BaseAVDCreate):
             kernel_image_path = None
             initramfs_image_path = None
 
+        if avd_spec.local_vendor_boot_image:
+            vendor_boot_image_path = create_common.FindVendorBootImage(
+                avd_spec.local_vendor_boot_image)
+
         if avd_spec.local_vendor_image:
             vendor_image_paths = cvd_utils.FindVendorImages(
                 avd_spec.local_vendor_image)
diff --git a/create/local_image_local_instance_test.py b/create/local_image_local_instance_test.py
index be6d312a..7475c88c 100644
--- a/create/local_image_local_instance_test.py
+++ b/create/local_image_local_instance_test.py
@@ -360,6 +360,7 @@ EOF"""
                 local_kernel_image=None,
                 local_system_image=None,
                 local_vendor_image=None,
+                local_vendor_boot_image=None,
                 local_tool_dirs=[cvd_dir])
 
             with self.assertRaisesRegex(
@@ -425,6 +426,7 @@ EOF"""
                 local_kernel_image=extra_image_dir,
                 local_system_image=extra_image_dir,
                 local_vendor_image=extra_image_dir,
+                local_vendor_boot_image=None,
                 local_tool_dirs=[])
 
             with mock.patch.dict("acloud.create.local_image_local_instance."
@@ -474,6 +476,7 @@ EOF"""
                 local_kernel_image=kernel_image_dir,
                 local_system_image=system_image_path,
                 local_vendor_image=None,
+                local_vendor_boot_image=None,
                 local_tool_dirs=[ota_tools_dir, cvd_dir])
 
             with mock.patch.dict("acloud.create.local_image_local_instance."
diff --git a/create/remote_image_local_instance.py b/create/remote_image_local_instance.py
index e8feaf39..445ab121 100644
--- a/create/remote_image_local_instance.py
+++ b/create/remote_image_local_instance.py
@@ -292,6 +292,10 @@ class RemoteImageLocalInstance(local_image_local_instance.LocalImageLocalInstanc
             ) = self.FindBootOrKernelImages(
                 os.path.abspath(avd_spec.local_kernel_image))
 
+        if avd_spec.local_vendor_boot_image:
+            vendor_boot_image_path = create_common.FindVendorBootImage(
+                avd_spec.local_vendor_boot_image)
+
         if avd_spec.local_vendor_image:
             vendor_image_paths = cvd_utils.FindVendorImages(
                 avd_spec.local_vendor_image)
diff --git a/create/remote_image_local_instance_test.py b/create/remote_image_local_instance_test.py
index 9a5ee483..8f9790b7 100644
--- a/create/remote_image_local_instance_test.py
+++ b/create/remote_image_local_instance_test.py
@@ -62,6 +62,7 @@ class RemoteImageLocalInstanceTest(driver_test_lib.BaseDriverTest):
         avd_spec.local_system_image = None
         avd_spec.local_kernel_image = None
         avd_spec.local_vendor_image = None
+        avd_spec.local_vendor_boot_image = None
         # raise errors.NoCuttlefishCommonInstalled
         self.Patch(setup_common, "PackageInstalled", return_value=False)
         self.assertRaises(errors.NoCuttlefishCommonInstalled,
diff --git a/internal/lib/cvd_utils.py b/internal/lib/cvd_utils.py
index 9d5a9b75..9943a1ce 100644
--- a/internal/lib/cvd_utils.py
+++ b/internal/lib/cvd_utils.py
@@ -313,14 +313,17 @@ def FindKernelImages(search_path):
 
 
 @utils.TimeExecute(function_description="Uploading local kernel images.")
-def _UploadKernelImages(ssh_obj, remote_image_dir, search_path):
+def _UploadKernelImages(ssh_obj, remote_image_dir, kernel_search_path,
+                        vendor_boot_search_path):
     """Find and upload kernel or boot images to a remote host or a GCE
     instance.
 
     Args:
         ssh_obj: An Ssh object.
         remote_image_dir: The remote image directory.
-        search_path: A path to an image file or an image directory.
+        kernel_search_path: A path to an image file or an image directory.
+        vendor_boot_search_path: A path to a vendor boot image file or an image
+                                 directory.
 
     Returns:
         A list of string pairs. Each pair consists of a launch_cvd option and a
@@ -334,7 +337,26 @@ def _UploadKernelImages(ssh_obj, remote_image_dir, search_path):
     ssh_obj.Run("mkdir -p " +
                 remote_path.join(remote_image_dir, _REMOTE_EXTRA_IMAGE_DIR))
 
-    kernel_image_path, initramfs_image_path = FindKernelImages(search_path)
+    # Find images
+    kernel_image_path = None
+    initramfs_image_path = None
+    boot_image_path = None
+    vendor_boot_image_path = None
+
+    if kernel_search_path:
+        kernel_image_path, initramfs_image_path = FindKernelImages(
+            kernel_search_path)
+        if not (kernel_image_path and initramfs_image_path):
+            boot_image_path, vendor_boot_image_path = FindBootImages(
+                kernel_search_path)
+
+    if vendor_boot_search_path:
+        vendor_boot_image_path = create_common.FindVendorBootImage(
+            vendor_boot_search_path)
+
+    # Upload
+    launch_cvd_args = []
+
     if kernel_image_path and initramfs_image_path:
         remote_kernel_image_path = remote_path.join(
             remote_image_dir, _REMOTE_KERNEL_IMAGE_PATH)
@@ -342,26 +364,29 @@ def _UploadKernelImages(ssh_obj, remote_image_dir, search_path):
             remote_image_dir, _REMOTE_INITRAMFS_IMAGE_PATH)
         ssh_obj.ScpPushFile(kernel_image_path, remote_kernel_image_path)
         ssh_obj.ScpPushFile(initramfs_image_path, remote_initramfs_image_path)
-        return [("-kernel_path", remote_kernel_image_path),
-                ("-initramfs_path", remote_initramfs_image_path)]
+        launch_cvd_args.append(("-kernel_path", remote_kernel_image_path))
+        launch_cvd_args.append(("-initramfs_path", remote_initramfs_image_path))
 
-    boot_image_path, vendor_boot_image_path = FindBootImages(search_path)
     if boot_image_path:
         remote_boot_image_path = remote_path.join(
             remote_image_dir, _REMOTE_BOOT_IMAGE_PATH)
         ssh_obj.ScpPushFile(boot_image_path, remote_boot_image_path)
-        launch_cvd_args = [("-boot_image", remote_boot_image_path)]
-        if vendor_boot_image_path:
-            remote_vendor_boot_image_path = remote_path.join(
-                remote_image_dir, _REMOTE_VENDOR_BOOT_IMAGE_PATH)
-            ssh_obj.ScpPushFile(vendor_boot_image_path,
-                                remote_vendor_boot_image_path)
-            launch_cvd_args.append(("-vendor_boot_image",
-                                    remote_vendor_boot_image_path))
-        return launch_cvd_args
+        launch_cvd_args.append(("-boot_image", remote_boot_image_path))
 
-    raise errors.GetLocalImageError(
-        f"{search_path} is not a boot image or a directory containing images.")
+    if vendor_boot_image_path:
+        remote_vendor_boot_image_path = remote_path.join(
+            remote_image_dir, _REMOTE_VENDOR_BOOT_IMAGE_PATH)
+        ssh_obj.ScpPushFile(vendor_boot_image_path,
+                            remote_vendor_boot_image_path)
+        launch_cvd_args.append(
+            ("-vendor_boot_image", remote_vendor_boot_image_path))
+
+    if not launch_cvd_args:
+        raise errors.GetLocalImageError(
+            f"{kernel_search_path}, {vendor_boot_search_path} is not a boot "
+            "image or a directory containing images.")
+
+    return launch_cvd_args
 
 
 def _FindSystemDlkmImage(search_path):
@@ -490,9 +515,11 @@ def UploadExtraImages(ssh_obj, remote_image_dir, avd_spec, target_files_dir):
         ValueError if target_files_dir is required but not specified.
     """
     extra_img_args = []
-    if avd_spec.local_kernel_image:
+    if avd_spec.local_kernel_image or avd_spec.local_vendor_boot_image:
         extra_img_args += _UploadKernelImages(ssh_obj, remote_image_dir,
-                                              avd_spec.local_kernel_image)
+                                              avd_spec.local_kernel_image,
+                                              avd_spec.local_vendor_boot_image)
+
 
     if AreTargetFilesRequired(avd_spec):
         if not target_files_dir:
diff --git a/internal/lib/cvd_utils_test.py b/internal/lib/cvd_utils_test.py
index 6455a770..d174cda2 100644
--- a/internal/lib/cvd_utils_test.py
+++ b/internal/lib/cvd_utils_test.py
@@ -156,7 +156,8 @@ class CvdUtilsTest(driver_test_lib.BaseDriverTest):
             mock_avd_spec = mock.Mock(local_kernel_image="boot.img",
                                       local_system_image=None,
                                       local_system_dlkm_image=None,
-                                      local_vendor_image=None)
+                                      local_vendor_image=None,
+                                      local_vendor_boot_image=None)
             args = cvd_utils.UploadExtraImages(mock_ssh, "dir", mock_avd_spec,
                                                None)
             self.assertEqual([("-boot_image", "dir/acloud_image/boot.img")],
@@ -188,7 +189,8 @@ class CvdUtilsTest(driver_test_lib.BaseDriverTest):
             mock_avd_spec = mock.Mock(local_kernel_image=kernel_image_path,
                                       local_system_image=None,
                                       local_system_dlkm_image=None,
-                                      local_vendor_image=None)
+                                      local_vendor_image=None,
+                                      local_vendor_boot_image=None)
             with self.assertRaises(errors.GetLocalImageError):
                 cvd_utils.UploadExtraImages(mock_ssh, "dir", mock_avd_spec,
                                             None)
@@ -220,6 +222,7 @@ class CvdUtilsTest(driver_test_lib.BaseDriverTest):
                                       local_system_image=extra_image_dir,
                                       local_system_dlkm_image=extra_image_dir,
                                       local_vendor_image=extra_image_dir,
+                                      local_vendor_boot_image=None,
                                       local_tool_dirs=[])
             self.CreateFile(
                 os.path.join(target_files_dir, "IMAGES", "boot.img"))
@@ -258,6 +261,44 @@ class CvdUtilsTest(driver_test_lib.BaseDriverTest):
         mock_ssh.ScpPushFile.assert_called_once_with(
             mock.ANY, "dir/acloud_image/vbmeta.img")
 
+
+    def testUploadVendorBootImages(self):
+        """Test UploadExtraImages."""
+        mock_ssh = mock.Mock()
+        with tempfile.TemporaryDirectory(prefix="cvd_utils") as image_dir:
+            vendor_boot_image_path = os.path.join(image_dir,
+                                                  "vendor_boot-debug_test.img")
+            self.CreateFile(vendor_boot_image_path)
+
+            mock_avd_spec = mock.Mock(
+                local_kernel_image=None,
+                local_system_image=None,
+                local_system_dlkm_image=None,
+                local_vendor_image=None,
+                local_vendor_boot_image=vendor_boot_image_path)
+
+            args = cvd_utils.UploadExtraImages(mock_ssh, "dir", mock_avd_spec,
+                                               None)
+            self.assertEqual(
+                [("-vendor_boot_image", "dir/acloud_image/vendor_boot.img")],
+                args)
+            mock_ssh.Run.assert_called_once()
+            mock_ssh.ScpPushFile.assert_called_once_with(
+                mock.ANY, "dir/acloud_image/vendor_boot.img")
+
+            mock_ssh.reset_mock()
+            self.CreateFile(os.path.join(image_dir, "vendor_boot.img"))
+            mock_avd_spec.local_vendor_boot_image = image_dir
+            args = cvd_utils.UploadExtraImages(mock_ssh, "dir", mock_avd_spec,
+                                               None)
+            self.assertEqual(
+                [("-vendor_boot_image", "dir/acloud_image/vendor_boot.img")],
+                args)
+            mock_ssh.Run.assert_called_once()
+            mock_ssh.ScpPushFile.assert_called_once_with(
+                mock.ANY, "dir/acloud_image/vendor_boot.img")
+
+
     def testCleanUpRemoteCvd(self):
         """Test CleanUpRemoteCvd."""
         mock_ssh = mock.Mock()