Create goldfish on remote host from local system image

This commit moves the pattern of system image name from cuttlefish and
goldfish modules to create_common, and makes
remote_host_gf_device_factory call the common function to find local
system image.

Test: acloud-dev create --avd-type goldfish -vv \
      --local-image ~/emu-extra-linux-system-images-8428355.zip \
      --emulator-zip ~/sdk-repo-linux-emulator-9714371.zip \
      --local-tool ~/otatools \
      --local-system-image ~/system.img \
      --host 192.168.9.2 --host-ssh-private-key-path ~/id_rsa
Bug: 245226952

Change-Id: I36ae9ef0c84cac672ac8d35930b8deded70e22c6

8 files changed, 61 insertions, 53 deletions

diff --git a/create/create_common.py b/create/create_common.py
index 4d654db1..00b5e23c 100644
--- a/create/create_common.py
+++ b/create/create_common.py
@@ -213,6 +213,11 @@ def FindBootImage(path, raise_error=True):
     return boot_image_path
 
 
+def FindSystemImage(path):
+    """Find a system image file in a given path."""
+    return FindLocalImage(path, _SYSTEM_IMAGE_NAME_PATTERN, raise_error=True)
+
+
 def DownloadRemoteArtifact(cfg, build_target, build_id, artifact, extract_path,
                            decompress=False):
     """Download remote artifact.
diff --git a/create/goldfish_local_image_local_instance.py b/create/goldfish_local_image_local_instance.py
index 75770a38..0ad37cad 100644
--- a/create/goldfish_local_image_local_instance.py
+++ b/create/goldfish_local_image_local_instance.py
@@ -59,7 +59,6 @@ logger = logging.getLogger(__name__)
 _EMULATOR_BIN_NAME = "emulator"
 _EMULATOR_BIN_DIR_NAMES = ("bin64", "qemu")
 _SDK_REPO_EMULATOR_DIR_NAME = "emulator"
-_SYSTEM_IMAGE_NAME_PATTERN = r"system\.img"
 _NON_MIXED_BACKUP_IMAGE_EXT = ".bak-non-mixed"
 _BUILD_PROP_FILE_NAME = "build.prop"
 # Timeout
@@ -449,8 +448,7 @@ class GoldfishLocalImageLocalInstance(base_avd_create.BaseAVDCreate):
             image_dir = self._FindImageDir(avd_spec.local_image_dir)
             mixed_image = goldfish_utils.MixWithSystemImage(
                 os.path.join(instance_dir, "mix_disk"), image_dir,
-                create_common.FindLocalImage(avd_spec.local_system_image,
-                                             _SYSTEM_IMAGE_NAME_PATTERN),
+                create_common.FindSystemImage(avd_spec.local_system_image),
                 ota_tools.FindOtaTools(ota_tools_search_paths))
 
             # TODO(b/142228085): Use -system instead of modifying image_dir.
diff --git a/create/local_image_local_instance.py b/create/local_image_local_instance.py
index a426ce99..e9b15d58 100644
--- a/create/local_image_local_instance.py
+++ b/create/local_image_local_instance.py
@@ -74,7 +74,6 @@ from acloud.setup import mkcert
 
 logger = logging.getLogger(__name__)
 
-_SYSTEM_IMAGE_NAME_PATTERN = r"system\.img"
 _SUPER_IMAGE_NAME = "super.img"
 _MIXED_SUPER_IMAGE_NAME = "mixed_super.img"
 _CMD_CVD_START = " start"
@@ -498,8 +497,8 @@ class LocalImageLocalInstance(base_avd_create.BaseAVDCreate):
             image_dir = cvd_utils.FindImageDir(image_dir)
             ota_tools_dir = os.path.abspath(
                 ota_tools.FindOtaToolsDir(tool_dirs))
-            system_image_path = create_common.FindLocalImage(
-                avd_spec.local_system_image, _SYSTEM_IMAGE_NAME_PATTERN)
+            system_image_path = create_common.FindSystemImage(
+                avd_spec.local_system_image)
         else:
             self._VerifyExtractedImgZip(image_dir)
             misc_info_path = None
diff --git a/create/remote_image_local_instance.py b/create/remote_image_local_instance.py
index 245ed178..4aaf4f57 100644
--- a/create/remote_image_local_instance.py
+++ b/create/remote_image_local_instance.py
@@ -53,7 +53,6 @@ _HOME_FOLDER = os.path.expanduser("~")
 # for the downloaded image artifacts.
 _REQUIRED_SPACE = 10
 
-_SYSTEM_IMAGE_NAME_PATTERN = r"system\.img"
 _SYSTEM_MIX_IMAGE_DIR = "mix_image_{build_id}"
 
 
@@ -267,8 +266,8 @@ class RemoteImageLocalInstance(local_image_local_instance.LocalImageLocalInstanc
                 except errors.GetCvdLocalHostPackageError:
                     logger.debug("fall back to downloaded cvd host binaries")
         if avd_spec.local_system_image:
-            system_image_path = create_common.FindLocalImage(
-                avd_spec.local_system_image, _SYSTEM_IMAGE_NAME_PATTERN)
+            system_image_path = create_common.FindSystemImage(
+                avd_spec.local_system_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 e470c913..18af6558 100644
--- a/create/remote_image_local_instance_test.py
+++ b/create/remote_image_local_instance_test.py
@@ -91,7 +91,8 @@ class RemoteImageLocalInstanceTest(driver_test_lib.BaseDriverTest):
         self.Patch(cvd_utils, "FindImageDir",
                    return_value="/mix_image_1234/IMAGES")
         self.Patch(ota_tools, "FindOtaToolsDir", return_value="/ota_tools_dir")
-        self.Patch(create_common, "FindLocalImage", return_value="/system_image_path")
+        self.Patch(create_common, "FindSystemImage",
+                   return_value="/system_image_path")
         self.Patch(self.RemoteImageLocalInstance, "_FindCvdHostBinaries",
                    side_effect=errors.GetCvdLocalHostPackageError("not found"))
         paths = self.RemoteImageLocalInstance.GetImageArtifactsPath(avd_spec)
diff --git a/public/actions/remote_host_gf_device_factory.py b/public/actions/remote_host_gf_device_factory.py
index 7f89b30b..c4b44c4d 100644
--- a/public/actions/remote_host_gf_device_factory.py
+++ b/public/actions/remote_host_gf_device_factory.py
@@ -48,8 +48,6 @@ _SDK_REPO_IMAGE_ZIP_NAME_FORMAT = ("sdk-repo-linux-system-images-"
 _EXTRA_IMAGE_ZIP_NAME_FORMAT = "emu-extra-linux-system-images-%(build_id)s.zip"
 _IMAGE_ZIP_NAME_FORMAT = "%(build_target)s-img-%(build_id)s.zip"
 _OTA_TOOLS_ZIP_NAME = "otatools.zip"
-_SYSTEM_IMAGE_NAME = "system.img"
-
 _EMULATOR_INFO_NAME = "emulator-info.txt"
 _EMULATOR_VERSION_PATTERN = re.compile(r"require\s+version-emulator="
                                        r"(?P<build_id>\w+)")
@@ -60,6 +58,7 @@ _SDK_REPO_EMULATOR_DIR_NAME = "emulator"
 # Files in temporary artifact directory.
 _DOWNLOAD_DIR_NAME = "download"
 _OTA_TOOLS_DIR_NAME = "ota_tools"
+_SYSTEM_IMAGE_NAME = "system.img"
 # Base directory of an instance.
 _REMOTE_INSTANCE_DIR_FORMAT = "acloud_gf_%d"
 # Relative paths in a base directory.
@@ -83,7 +82,7 @@ _MISSING_EMULATOR_MSG = ("No emulator zip. Specify "
 ArtifactPaths = collections.namedtuple(
     "ArtifactPaths",
     ["image_zip", "emulator_zip", "ota_tools_dir",
-     "system_image_zip", "boot_image"])
+     "system_image", "boot_image"])
 
 RemotePaths = collections.namedtuple(
     "RemotePaths",
@@ -339,7 +338,12 @@ class RemoteHostGoldfishDeviceFactory(base_device_factory.BaseDeviceFactory):
         if not emu_zip_path:
             raise errors.GetSdkRepoPackageError(_MISSING_EMULATOR_MSG)
 
-        system_image_zip_path = self._RetrieveSystemImageZip()
+        # System image.
+        if self._avd_spec.local_system_image:
+            system_image_path = create_common.FindSystemImage(
+                self._avd_spec.local_system_image)
+        else:
+            system_image_path = self._RetrieveSystemImage()
 
         # Boot image.
         if self._avd_spec.local_kernel_image:
@@ -350,7 +354,7 @@ class RemoteHostGoldfishDeviceFactory(base_device_factory.BaseDeviceFactory):
 
         # OTA tools.
         ota_tools_dir = None
-        if system_image_zip_path or boot_image_path:
+        if system_image_path or boot_image_path:
             if self._avd_spec.image_source == constants.IMAGE_SRC_REMOTE:
                 ota_tools_dir = self._RetrieveOtaTools()
             else:
@@ -361,7 +365,7 @@ class RemoteHostGoldfishDeviceFactory(base_device_factory.BaseDeviceFactory):
                         constants.ENV_ANDROID_HOST_OUT))
 
         return ArtifactPaths(image_zip_path, emu_zip_path, ota_tools_dir,
-                             system_image_zip_path, boot_image_path)
+                             system_image_path, boot_image_path)
 
     def _RetrieveDeviceImageZip(self):
         """Retrieve device image zip from cache or Android Build API.
@@ -416,23 +420,28 @@ class RemoteHostGoldfishDeviceFactory(base_device_factory.BaseDeviceFactory):
         return self._RetrieveArtifact(emu_build_target, emu_build_id,
                                       emu_zip_name)
 
-    def _RetrieveSystemImageZip(self):
-        """Retrieve system image zip if system build info is not empty.
+    def _RetrieveSystemImage(self):
+        """Retrieve and unzip system image if system build info is not empty.
 
         Returns:
-            The path to the system image zip in download_dir.
+            The path to the temporary system image.
             None if the system build info is empty.
         """
         build_id = self._avd_spec.system_build_info.get(constants.BUILD_ID)
         build_target = self._avd_spec.system_build_info.get(
             constants.BUILD_TARGET)
-        if build_id and build_target:
-            image_zip_name = _IMAGE_ZIP_NAME_FORMAT % {
-                "build_target": build_target.split("-", 1)[0],
-                "build_id": build_id}
-            return self._RetrieveArtifact(build_target, build_id,
-                                          image_zip_name)
-        return None
+        if not build_id or not build_target:
+            return None
+        image_zip_name = _IMAGE_ZIP_NAME_FORMAT % {
+            "build_target": build_target.split("-", 1)[0],
+            "build_id": build_id}
+        image_zip_path = self._RetrieveArtifact(build_target, build_id,
+                                                image_zip_name)
+        logger.debug("Unzip %s from %s to %s.",
+                     _SYSTEM_IMAGE_NAME, image_zip_path, self._artifact_dir)
+        with zipfile.ZipFile(image_zip_path, "r") as zip_file:
+            zip_file.extract(_SYSTEM_IMAGE_NAME, self._artifact_dir)
+        return os.path.join(self._artifact_dir, _SYSTEM_IMAGE_NAME)
 
     def _RetrieveBootImage(self):
         """Retrieve boot image if boot build info is not empty.
@@ -499,7 +508,7 @@ class RemoteHostGoldfishDeviceFactory(base_device_factory.BaseDeviceFactory):
                            zip_path, " ".join(entries))
             return ""
 
-    def _UploadArtifacts(self, artifacts_paths):
+    def _UploadArtifacts(self, artifact_paths):
         """Process and upload all images and tools to the remote host.
 
         Args:
@@ -509,33 +518,33 @@ class RemoteHostGoldfishDeviceFactory(base_device_factory.BaseDeviceFactory):
             An object of RemotePaths.
         """
         remote_emulator_dir, remote_image_dir = self._UploadDeviceImages(
-            artifacts_paths.emulator_zip, artifacts_paths.image_zip)
+            artifact_paths.emulator_zip, artifact_paths.image_zip)
 
         remote_kernel_path = None
         remote_ramdisk_path = None
 
-        if artifacts_paths.boot_image or artifacts_paths.system_image_zip:
+        if artifact_paths.boot_image or artifact_paths.system_image:
             with tempfile.TemporaryDirectory("host_gf") as temp_dir:
-                ota = ota_tools.OtaTools(artifacts_paths.ota_tools_dir)
+                ota = ota_tools.OtaTools(artifact_paths.ota_tools_dir)
 
                 image_dir = os.path.join(temp_dir, "images")
-                logger.debug("Unzip %s.", artifacts_paths.image_zip)
-                with zipfile.ZipFile(artifacts_paths.image_zip,
+                logger.debug("Unzip %s.", artifact_paths.image_zip)
+                with zipfile.ZipFile(artifact_paths.image_zip,
                                      "r") as zip_file:
                     zip_file.extractall(image_dir)
                 image_dir = os.path.join(
                     image_dir,
-                    self._GetSubdirNameInZip(artifacts_paths.image_zip))
+                    self._GetSubdirNameInZip(artifact_paths.image_zip))
 
-                if artifacts_paths.system_image_zip:
+                if artifact_paths.system_image:
                     self._MixAndUploadDiskImage(
                         remote_image_dir, image_dir,
-                        artifacts_paths.system_image_zip, ota)
+                        artifact_paths.system_image, ota)
 
-                if artifacts_paths.boot_image:
+                if artifact_paths.boot_image:
                     remote_kernel_path, remote_ramdisk_path = (
                         self._MixAndUploadKernelImages(
-                            image_dir, artifacts_paths.boot_image, ota))
+                            image_dir, artifact_paths.boot_image, ota))
 
         return RemotePaths(remote_image_dir, remote_emulator_dir,
                            remote_kernel_path, remote_ramdisk_path)
@@ -550,8 +559,9 @@ class RemoteHostGoldfishDeviceFactory(base_device_factory.BaseDeviceFactory):
         Returns:
             Boolean, whether a mixed disk image is required.
         """
-        return (self._avd_spec.system_build_info.get(constants.BUILD_ID) and
-                self._avd_spec.system_build_info.get(constants.BUILD_TARGET))
+        return self._avd_spec.local_system_image or (
+            self._avd_spec.system_build_info.get(constants.BUILD_ID) and
+            self._avd_spec.system_build_info.get(constants.BUILD_TARGET))
 
     @utils.TimeExecute(
         function_description="Processing and uploading tools and images")
@@ -597,30 +607,22 @@ class RemoteHostGoldfishDeviceFactory(base_device_factory.BaseDeviceFactory):
         return remote_emulator_subdir, remote_image_subdir
 
     def _MixAndUploadDiskImage(self, remote_image_dir, image_dir,
-                               system_image_zip_path, ota):
+                               system_image_path, ota):
         """Mix emulator images with a system image and upload them.
 
         Args:
             remote_image_dir: The remote directory where the mixed disk image
                               is uploaded.
             image_dir: The directory containing emulator images.
-            system_image_zip_path: The path to the zip containing the system
-                                   image.
+            system_image_path: The path to the system image.
             ota: An instance of ota_tools.OtaTools.
 
         Returns:
             The remote path to the mixed disk image.
         """
         with tempfile.TemporaryDirectory("host_gf_disk") as temp_dir:
-            logger.debug("Unzip %s.", system_image_zip_path)
-            with zipfile.ZipFile(system_image_zip_path, "r") as zip_file:
-                zip_file.extract(_SYSTEM_IMAGE_NAME, temp_dir)
-
             mixed_image = goldfish_utils.MixWithSystemImage(
-                os.path.join(temp_dir, "mix_disk"),
-                image_dir,
-                os.path.join(temp_dir, _SYSTEM_IMAGE_NAME),
-                ota)
+                temp_dir, image_dir, system_image_path, ota)
 
             # TODO(b/142228085): Use -system instead of overwriting the file.
             remote_disk_image_path = os.path.join(
diff --git a/public/actions/remote_host_gf_device_factory_test.py b/public/actions/remote_host_gf_device_factory_test.py
index 0669e02a..f0d5a9cd 100644
--- a/public/actions/remote_host_gf_device_factory_test.py
+++ b/public/actions/remote_host_gf_device_factory_test.py
@@ -65,6 +65,7 @@ class RemoteHostGoldfishDeviceFactoryTest(driver_test_lib.BaseDriverTest):
         "boot_build_info": {},
         "local_image_artifact": None,
         "local_kernel_image": None,
+        "local_system_image": None,
         "local_tool_dirs": [],
         "base_instance_num": None,
         "boot_timeout_secs": None,
@@ -319,11 +320,14 @@ class RemoteHostGoldfishDeviceFactoryTest(driver_test_lib.BaseDriverTest):
             self._CreateSdkRepoZip(image_zip_path)
             boot_image_path = os.path.join(temp_dir, "boot.img")
             self.CreateFile(boot_image_path, b"ANDROID!")
+            system_image_path = os.path.join(temp_dir, "system.img")
+            self.CreateFile(system_image_path)
             self._mock_avd_spec.emulator_zip = emulator_zip_path
             self._mock_avd_spec.image_source = constants.IMAGE_SRC_LOCAL
             self._mock_avd_spec.remote_image = {}
             self._mock_avd_spec.local_image_artifact = image_zip_path
             self._mock_avd_spec.local_kernel_image = boot_image_path
+            self._mock_avd_spec.local_system_image = system_image_path
             self._mock_avd_spec.local_tool_dirs.append("/otatools")
             mock_gf_utils.ConvertAvdSpecToArgs.return_value = ["-gpu", "auto"]
             mock_gf_utils.MixWithBootImage.return_value = (
@@ -336,6 +340,7 @@ class RemoteHostGoldfishDeviceFactoryTest(driver_test_lib.BaseDriverTest):
             factory.CreateInstance()
 
             mock_gf_utils.MixWithBootImage.assert_called_once()
+            mock_gf_utils.MixWithSystemImage.assert_called_once()
             mock_ota_tools.FindOtaToolsDir.assert_called_once()
             self.assertEqual("/otatools",
                              mock_ota_tools.FindOtaToolsDir.call_args[0][0][0])
diff --git a/public/actions/remote_instance_cf_device_factory.py b/public/actions/remote_instance_cf_device_factory.py
index e16d070e..072cb39e 100644
--- a/public/actions/remote_instance_cf_device_factory.py
+++ b/public/actions/remote_instance_cf_device_factory.py
@@ -30,7 +30,6 @@ from acloud.pull import pull
 
 logger = logging.getLogger(__name__)
 _SCREEN_CONSOLE_COMMAND = "screen ~/cuttlefish_runtime/console"
-_SYSTEM_IMAGE_NAME_PATTERN = r"system\.img"
 _MIXED_SUPER_IMAGE_NAME = "mixed_super.img"
 
 
@@ -183,8 +182,8 @@ class RemoteInstanceDeviceFactory(gce_device_factory.GCEDeviceFactory):
         odm_dlkm_image_path=None
 
         if avd_spec.local_system_image:
-            system_image_path = create_common.FindLocalImage(
-                avd_spec.local_system_image, _SYSTEM_IMAGE_NAME_PATTERN)
+            system_image_path = create_common.FindSystemImage(
+                avd_spec.local_system_image)
 
         if avd_spec.local_vendor_image:
             vendor_image_paths = cvd_utils.FindVendorImages(