xbuddy: Do not rename update payloads to update.gz anymore

This reverts commit b9de5766c2266f8c3d1b689346961cc945a707ad.

This allows us to have more update payload types and we don't have to
put delta update payloads in 'nton' directory anymore.

There is no change in the behavior of cros flash.

This CL basically removes the GS to Local name mapping and instead
relies on the file names provided by the artifacts themselves.

Also removes an unneeded logic for checking the cached version of files
which xbuddy is already handling fine.

BUG=chromium:1008058
TEST=cros flash --debug <dut-ip> xbuddy://remote/reef-release/R86-13421.67.0/
TEST=cros flash --debug <dut-ip> xbuddy://remote/reef-release/R86-13421.67.0/test
TEST=cros flash --debug file://some-file xbuddy://remote/reef-release/R86-13421.67.0/
TEST=cros flash --debug <dut-ip> latest

Change-Id: If9c2850531c4ed3e9f3ee501d22fcd39618a058d
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/2508870
Tested-by: Amin Hassani <ahassani@chromium.org>
Auto-Submit: Amin Hassani <ahassani@chromium.org>
Reviewed-by: Achuith Bhandarkar <achuith@chromium.org>
Commit-Queue: Amin Hassani <ahassani@chromium.org>
diff --git a/cli/flash.py b/cli/flash.py
index bfcf0ee..2028516 100644
--- a/cli/flash.py
+++ b/cli/flash.py
@@ -400,8 +400,8 @@
       logging.info('You can find the log files and/or payloads in %s',
                    self.tempdir)
 
-  def GetPayloadDir(self, device):
-    """Get directory of payload for update.
+  def GetPayloadPaths(self, device):
+    """Get directory of payload and rootfs payload file name for update.
 
     This method is used to obtain the directory of payload for cros-flash. The
     given path 'self.image' is passed in when initializing RemoteDeviceUpdater.
@@ -427,19 +427,21 @@
       device: A ChromiumOSDevice object.
 
     Returns:
-      A string payload_dir, that represents the payload directory.
+      A string tuple (payload_dir, rootfs_filename). payload_dir is the
+      directory where the update payloads are located. rootfs_filename is the
+      name of the rootfs update payload (sometimes update.gz).
     """
+    rootfs_filename = auto_updater_transfer.ROOTFS_FILENAME
+
     if os.path.isdir(self.image):
       # The given path is a directory.
       logging.info('Using provided payloads in %s', self.image)
-      return self.image
+      return self.image, rootfs_filename
 
-    image_path = None
-    if os.path.isfile(self.image):
-      # The given path is an image.
-      image_path = self.image
-      payload_dir = self.tempdir
-    else:
+    image_path = self.image
+    payload_dir = self.tempdir
+
+    if not os.path.isfile(self.image):
       # Assuming it is an xbuddy path.
       self.board = cros_build_lib.GetBoard(
           device_board=device.board or GetDefaultBoard(),
@@ -461,8 +463,9 @@
         translated_path, _ = ds_wrapper.GetImagePathWithXbuddy(
             os.path.join(self.image, artifact_info.FULL_PAYLOAD),
             self.board, self.version, silent=True)
-        payload_dir = os.path.dirname(
-            ds_wrapper.TranslatedPathToLocalPath(translated_path))
+        local_path = ds_wrapper.TranslatedPathToLocalPath(translated_path)
+        payload_dir, rootfs_filename = os.path.split(local_path)
+
         ds_wrapper.GetImagePathWithXbuddy(
             os.path.join(self.image, artifact_info.STATEFUL_PAYLOAD),
             self.board, self.version, silent=True)
@@ -482,8 +485,7 @@
                        image_path, payload_dir)
 
     # Generate rootfs and stateful update payloads if they do not exist.
-    payload_path = os.path.join(payload_dir,
-                                auto_updater_transfer.ROOTFS_FILENAME)
+    payload_path = os.path.join(payload_dir, rootfs_filename)
     if not os.path.exists(payload_path):
       paygen_payload_lib.GenerateUpdatePayload(
           image_path, payload_path, src_image=self.src_image_to_delta)
@@ -491,7 +493,7 @@
         payload_dir, auto_updater_transfer.STATEFUL_FILENAME)):
       paygen_stateful_payload_lib.GenerateStatefulPayload(image_path,
                                                           payload_dir)
-    return payload_dir
+    return payload_dir, rootfs_filename
 
   def Run(self):
     """Perform remote device update.
@@ -510,14 +512,14 @@
         try:
           # Get payload directory
           logging.notice('Staging payloads...')
-          payload_dir = self.GetPayloadDir(device)
+          payload_dir, rootfs_filename = self.GetPayloadPaths(device)
 
           # Do auto-update
           chromeos_AU = auto_updater.ChromiumOSUpdater(
               device=device,
               build_name=None,
               payload_dir=payload_dir,
-              payload_filename=auto_updater_transfer.ROOTFS_FILENAME,
+              payload_filename=rootfs_filename,
               tempdir=self.tempdir,
               do_rootfs_update=self.do_rootfs_update,
               do_stateful_update=self.do_stateful_update,
diff --git a/lib/xbuddy/build_artifact.py b/lib/xbuddy/build_artifact.py
index cbd7322..9830e56 100644
--- a/lib/xbuddy/build_artifact.py
+++ b/lib/xbuddy/build_artifact.py
@@ -25,8 +25,6 @@
 # We do a number of things with args/kwargs arguments that confuse pylint.
 # pylint: disable=docstring-misnamed-args
 
-AU_NTON_DIR = 'au_nton'
-
 ############ Actual filenames of artifacts in Google Storage ############
 
 AU_SUITE_FILE = 'au_control.tar.bz2'
@@ -236,6 +234,12 @@
     self._Log('ArtifactStaged() -> yes, %s is staged.', self)
     return True
 
+  def StagedFiles(self):
+    """Returns the installed/staged files for this artifact."""
+    with open(os.path.join(self.install_dir, self.marker_name)) as f:
+      return [line.strip() for line in f]
+
+
   def _MarkArtifactStaged(self):
     """Marks the artifact as staged."""
     with open(os.path.join(self.install_dir, self.marker_name), 'w') as f:
@@ -353,8 +357,6 @@
           # Save the exception to a file for downloader.IsStaged to retrieve.
           self._SaveException(e)
           raise e
-      else:
-        self._Log('%s is already staged.', self)
 
   def __str__(self):
     """String representation for the download."""
@@ -387,31 +389,6 @@
                                          name) for name in self.name]
 
 
-class AUTestPayload(MultiArtifact):
-  """Wrapper for AUTest delta payloads which need additional setup."""
-
-  def _Setup(self):
-    super(AUTestPayload, self)._Setup()
-
-    # Rename to update.gz.
-    # TODO(crbug.com/1008058): Change the devserver such that this renaming is
-    # not needed anymore.
-    for name in self.name:
-      dest_name = (devserver_constants.UPDATE_METADATA_FILE
-                   if name.endswith('.json')
-                   else devserver_constants.UPDATE_FILE)
-
-      install_path = os.path.join(self.install_dir, self.install_subdir, name)
-      new_install_path = os.path.join(self.install_dir, self.install_subdir,
-                                      dest_name)
-      self._Log('moving %s to %s', install_path, new_install_path)
-      shutil.move(install_path, new_install_path)
-
-      # Reflect the rename in the list of installed files.
-      self.installed_files.remove(install_path)
-      self.installed_files.append(new_install_path)
-
-
 class BundledArtifact(Artifact):
   """A single build artifact bundle e.g. zip file or tar file."""
 
@@ -593,13 +570,12 @@
   chromeos_artifact_map.setdefault(tag, []).append(artifact)
 
 
-_AddCrOSArtifact(artifact_info.FULL_PAYLOAD, AUTestPayload,
+_AddCrOSArtifact(artifact_info.FULL_PAYLOAD, MultiArtifact,
                  r'chromeos_.*_full_dev.*bin(\.json)?\Z', is_regex_name=True,
                  alt_name=[u'chromeos_{build}_{board}_dev.bin',
                            u'chromeos_{build}_{board}_dev.bin.json'])
-_AddCrOSArtifact(artifact_info.DELTA_PAYLOAD, AUTestPayload,
-                 r'chromeos_.*_delta_dev.*bin(\.json)?\Z', is_regex_name=True,
-                 install_subdir=AU_NTON_DIR)
+_AddCrOSArtifact(artifact_info.DELTA_PAYLOAD, MultiArtifact,
+                 r'chromeos_.*_delta_dev.*bin(\.json)?\Z', is_regex_name=True)
 _AddCrOSArtifact(artifact_info.STATEFUL_PAYLOAD, Artifact,
                  devserver_constants.STATEFUL_FILE)
 _AddCrOSArtifact(artifact_info.BASE_IMAGE, BundledArtifact, IMAGE_FILE,
diff --git a/lib/xbuddy/xbuddy.py b/lib/xbuddy/xbuddy.py
index 58d2615..bf20cb5 100644
--- a/lib/xbuddy/xbuddy.py
+++ b/lib/xbuddy/xbuddy.py
@@ -77,7 +77,6 @@
     RECOVERY,
     FACTORY_SHIM,
     FULL,
-    DELTA,
     STATEFUL,
     ANY,
 ]
@@ -108,17 +107,6 @@
     AUTOTEST,
 ]
 
-GS_FILE_NAMES = [
-    devserver_constants.TEST_IMAGE_FILE,
-    devserver_constants.BASE_IMAGE_FILE,
-    devserver_constants.RECOVERY_IMAGE_FILE,
-    devserver_constants.SIGNED_IMAGE_FILE,
-    devserver_constants.FACTORY_SHIM_IMAGE_FILE,
-    devserver_constants.UPDATE_FILE,
-    devserver_constants.STATEFUL_FILE,
-    devserver_constants.AUTOTEST_DIR,
-]
-
 ARTIFACTS = [
     artifact_info.TEST_IMAGE,
     artifact_info.BASE_IMAGE,
@@ -131,7 +119,6 @@
     artifact_info.AUTOTEST,
 ]
 
-GS_ALIAS_TO_FILENAME = dict(zip(GS_ALIASES, GS_FILE_NAMES))
 GS_ALIAS_TO_ARTIFACT = dict(zip(GS_ALIASES, ARTIFACTS))
 
 LATEST_OFFICIAL = 'latest-official'
@@ -708,6 +695,10 @@
   def _Download(self, gs_url, artifacts, build_id):
     """Download the artifacts from the given gs_url.
 
+    Returns:
+      A list containing lists of downloaded files for each artifact.
+      e.g.: artifacts = ['a', 'b'] -> return [['f1', 'f2'], ['f3']]
+
     Raises:
       build_artifact.ArtifactDownloadError: If we failed to download the
                                             artifact.
@@ -720,6 +711,8 @@
       factory = build_artifact.ChromeOSArtifactFactory(
           dl.GetBuildDir(), artifacts, [], dl.GetBuild())
       dl.Download(factory)
+      downloaded_artifacts = factory.RequiredArtifacts()
+      return [x.StagedFiles() for x in downloaded_artifacts]
     finally:
       with XBuddy._staging_thread_count_lock:
         XBuddy._staging_thread_count -= 1
@@ -808,25 +801,21 @@
       channel: The channel for the image. If none, it tries to guess it in
         order of stability.
 
+    Returns:
+      The list of files downloaded for the given image_type.
+
     Raises:
       build_artifact.ArtifactDownloadError: If we failed to download the
         artifact.
     """
-    # Stage image if not found in cache.
-    file_name = GS_ALIAS_TO_FILENAME[image_type]
-    file_loc = os.path.join(self.static_dir, build_id, file_name)
-    cached = os.path.exists(file_loc)
-
-    if not cached:
-      artifact = GS_ALIAS_TO_ARTIFACT[image_type]
-      if image_type == SIGNED:
-        gs_url = self._TranslateSignedGSUrl(build_id, channel=channel)
-      else:
-        image_dir = XBuddy._ResolveImageDir(image_dir)
-        gs_url = os.path.join(image_dir, build_id)
-      self._Download(gs_url, [artifact], build_id)
+    artifact = GS_ALIAS_TO_ARTIFACT[image_type]
+    if image_type == SIGNED:
+      gs_url = self._TranslateSignedGSUrl(build_id, channel=channel)
     else:
-      _Log('Image already cached.')
+      image_dir = XBuddy._ResolveImageDir(image_dir)
+      gs_url = os.path.join(image_dir, build_id)
+
+    return self._Download(gs_url, [artifact], build_id)[0]
 
   def _GetArtifact(self, path_list, board=None, version=None, image_dir=None):
     """Interpret an xBuddy path and return directory/file_name to resource.
@@ -878,7 +867,6 @@
       if not os.path.exists(artifact_path):
         raise XBuddyException('Local %s artifact not in static_dir at %s' %
                               (image_type, artifact_path))
-
     else:
       # Get a remote image.
       if image_type not in GS_ALIASES:
@@ -887,9 +875,8 @@
       build_id, channel = self._ResolveVersionToBuildIdAndChannel(
           board, suffix, version, image_dir=image_dir)
       _Log('Resolved version %s to %s.', version, build_id)
-      file_name = GS_ALIAS_TO_FILENAME[image_type]
-      self._GetFromGS(build_id, image_type, image_dir=image_dir,
-                      channel=channel)
+      file_name = self._GetFromGS(build_id, image_type, image_dir=image_dir,
+                                  channel=channel)[0]
 
     return build_id, file_name