Add the ability to quick-provision without lab cache servers

Running m2n AU tests locally is very painful for first-time users since
it requires extensive SSH configuration to the lab cache servers. SSH
access to lab cache servers from your workstation is also in a weird
place in terms of support right now, with some sources recommending
against the steps needed to set it up.

To provide a much easier alternative, I made some edits to
provisioner.py to allow using the public throwaway cellular bucket as
the source for the quick-provision artifacts, and applied these to
au_Basic. This way, no special setup is needed to run m2n tests for the
first time (except for possibly copying your GS .boto key into your
chroot) - no ssh access, no shadow_config.ini, no job_repo_urls.

I still need to update the other tests besides au_Basic with the changes
for running without job_repo_url, I will address those in follow-up CLs.

BUG=b:210547817
TEST=test_that --autotest_dir ~/chromiumos/src/third_party/autotest/files chromeos6-row17-rack22-host13 autoupdate_Basic.m2n.full --args="running_at_desk=True build=R104-14852.0.0"

Change-Id: I081376d76e9a08ef6d39a1c7f8ba2eb871932ff1
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/autotest/+/3658060
Tested-by: Kyle Shimabukuro <kyleshima@chromium.org>
Auto-Submit: Kyle Shimabukuro <kyleshima@chromium.org>
Commit-Queue: Kyle Shimabukuro <kyleshima@chromium.org>
Reviewed-by: Jae Hoon Kim <kimjae@chromium.org>

5 files changed, 123 insertions, 44 deletions

diff --git a/server/cros/provisioner.py b/server/cros/provisioner.py
index 916a575501..b209c5797c 100644
--- a/server/cros/provisioner.py
+++ b/server/cros/provisioner.py
@@ -218,7 +218,8 @@ class ChromiumOSProvisioner(object):
                  host=None,
                  interactive=True,
                  is_release_bucket=None,
-                 is_servohost=False):
+                 is_servohost=False,
+                 public_bucket=False):
         """Initializes the object.
 
         @param update_url: The URL we want the update to use.
@@ -227,6 +228,9 @@ class ChromiumOSProvisioner(object):
         @param is_release_bucket: If True, use release bucket
             gs://chromeos-releases.
         @param is_servohost: Bool whether the update target is a servohost.
+        @param public_bucket: True to copy payloads to a public throwaway GS
+            bucket. This avoids using a lab cache server, so local test runs
+            can provision without any special setup.
         """
         self.update_url = update_url
         self.host = host
@@ -234,6 +238,7 @@ class ChromiumOSProvisioner(object):
         self.update_version = _url_to_version(update_url)
         self._is_release_bucket = is_release_bucket
         self._is_servohost = is_servohost
+        self._public_bucket = public_bucket
 
     def _run(self, cmd, *args, **kwargs):
         """Abbreviated form of self.host.run(...)"""
@@ -404,6 +409,20 @@ class ChromiumOSProvisioner(object):
                         'gs_cache': False
                 })
 
+    def _quick_provision_with_public_bucket(self, provision_command,
+                                            image_name):
+        """Run quick_provision using public GS bucket.
+
+        @param provision_command: The path of quick_provision command.
+        @param image_name: The image to be installed.
+        """
+        logging.info('Try quick provision with public bucket.')
+
+        bucket_url = self.update_url[:self.update_url.find(image_name) - 1]
+        command = '%s --noreboot %s %s' % (provision_command, image_name,
+                                           bucket_url)
+        self._run(command)
+
     def _install_update(self):
         """Install an updating using the `quick-provision` script.
 
@@ -415,20 +434,27 @@ class ChromiumOSProvisioner(object):
         logging.info('Installing image at %s onto %s', self.update_url,
                      self.host.hostname)
         server_name = six.moves.urllib.parse.urlparse(self.update_url)[1]
-        image_name = url_to_image_name(self.update_url)
+        if self._public_bucket:
+            image_name = self.update_url.partition('provision/')[2]
+        else:
+            image_name = url_to_image_name(self.update_url)
 
         logging.info('Installing image using quick-provision.')
         provision_command = self._get_remote_script(_QUICK_PROVISION_SCRIPT)
         try:
-            try:
-                self._quick_provision_with_gs_cache(provision_command,
-                                                    server_name, image_name)
-            except Exception as e:
-                logging.error(
-                        'Failed to quick-provision with gscache with '
-                        'error %s', e)
-                self._quick_provision_with_devserver(provision_command,
-                                                     server_name, image_name)
+            if self._public_bucket:
+                self._quick_provision_with_public_bucket(
+                        provision_command, image_name)
+            else:
+                try:
+                    self._quick_provision_with_gs_cache(
+                            provision_command, server_name, image_name)
+                except Exception as e:
+                    logging.error(
+                            'Failed to quick-provision with gscache with '
+                            'error %s', e)
+                    self._quick_provision_with_devserver(
+                            provision_command, server_name, image_name)
 
             self._set_target_version()
             return kernel_utils.verify_kernel_state_after_update(self.host)
@@ -506,9 +532,11 @@ class ChromiumOSProvisioner(object):
             `image_name` is the name of the image installed, and
             `attributes` is new attributes to be applied to the DUT.
         """
-        server_name = dev_server.get_resolved_hostname(self.update_url)
-        metrics.Counter(_metric_name('install')).increment(
-                fields={'devserver': server_name})
+        server_name = ""
+        if not self._public_bucket:
+            server_name = dev_server.get_resolved_hostname(self.update_url)
+            metrics.Counter(_metric_name('install')).increment(
+                    fields={'devserver': server_name})
 
         try:
             self._prepare_host()
diff --git a/server/cros/update_engine/update_engine_test.py b/server/cros/update_engine/update_engine_test.py
index 6171d8bc43..afc3fb5677 100644
--- a/server/cros/update_engine/update_engine_test.py
+++ b/server/cros/update_engine/update_engine_test.py
@@ -7,6 +7,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import base64
 import json
 import logging
 import os
@@ -20,6 +21,7 @@ import six.moves.urllib.parse
 from datetime import datetime, timedelta
 from xml.etree import ElementTree
 
+from autotest_lib.client.common_lib import autotemp
 from autotest_lib.client.common_lib import autotest_enum
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.common_lib import global_config
@@ -36,6 +38,9 @@ from autotest_lib.server.cros import gsutil_wrapper
 from autotest_lib.server.cros.dynamic_suite import tools
 from autotest_lib.utils.frozen_chromite.lib import auto_updater
 from autotest_lib.utils.frozen_chromite.lib import auto_updater_transfer
+from autotest_lib.utils.frozen_chromite.lib import constants as chromite_constants
+from autotest_lib.utils.frozen_chromite.lib import gob_util
+from autotest_lib.utils.frozen_chromite.lib import osutils
 from autotest_lib.utils.frozen_chromite.lib import remote_access
 from autotest_lib.utils.frozen_chromite.lib import retry_util
 
@@ -76,6 +81,8 @@ class UpdateEngineTest(test.test, update_engine_util.UpdateEngineUtil):
     _CORRUPT_STATEFUL_PATH = '/mnt/stateful_partition/.corrupt_stateful'
 
     _STATEFUL_ARCHIVE_NAME = 'stateful.tgz'
+    _KERNEL_ARCHIVE_NAME = 'full_dev_part_KERN.bin.gz'
+    _ROOTFS_ARCHIVE_NAME = 'full_dev_part_ROOT.bin.gz'
 
     _PAYLOAD_TYPE = autotest_enum.AutotestEnum('CROS', 'DLC', 'MINIOS')
 
@@ -980,6 +987,58 @@ class UpdateEngineTest(test.test, update_engine_util.UpdateEngineUtil):
         logging.info('Public stateful URL: %s', url)
         return url
 
+    def _get_provision_url_on_public_bucket(self, release_path):
+        """
+        Copy the necessary artifacts for quick-provision to the public bucket
+        and return the URL pointing to them.
+
+        This is to enable local runs of tests that need to provision a source
+        version on the DUT (such as m2n tests) without requiring lab cache
+        server SSH access.
+
+        @param release_path: path to the build artifacts in
+            gs://chromeos-releases. Ex: dev-channel/asurada/14515.0.0. The
+            output of _get_latest_serving_stable_build matches this format.
+
+        """
+        # We have a flat directory structure in the public directory. Therefore
+        # we need to disambiguate the path to the provision artifacts.
+        new_gs_dir = os.path.join(self._CELLULAR_BUCKET, 'provision',
+                                  release_path)
+        src_gs_dir = os.path.join('gs://chromeos-releases', release_path)
+        provision_artifacts = [
+                self._STATEFUL_ARCHIVE_NAME, self._ROOTFS_ARCHIVE_NAME,
+                self._KERNEL_ARCHIVE_NAME
+        ]
+
+        for file in provision_artifacts:
+            src_url = os.path.join(src_gs_dir, file)
+            dst_url = os.path.join(new_gs_dir, file)
+            utils.run([
+                    'gsutil', 'cp', '-n', '-a', 'public-read', src_url, dst_url
+            ])
+
+        public_url = new_gs_dir.replace('gs://',
+                                        'https://storage.googleapis.com/')
+
+        return public_url
+
+    def _copy_quick_provision_to_dut(self):
+        """ Copies the quick-provision script to the DUT from googlesource."""
+        tmp = autotemp.tempdir(unique_id='m2n')
+        src = os.path.join(tmp.name, 'quick-provision')
+        dst = '/usr/local/bin/quick-provision'
+        logging.info('Downloading quick-provision from googlesource')
+        qp_url_path = '%s/+/%s/%s?format=text' % (
+                'chromiumos/platform/dev-util', 'refs/heads/main',
+                'quick-provision/quick-provision')
+        contents_b64 = gob_util.FetchUrl(chromite_constants.EXTERNAL_GOB_HOST,
+                                         qp_url_path)
+        osutils.WriteFile(src, base64.b64decode(contents_b64).decode('utf-8'))
+        self._host.send_file(src, dst)
+        self._run(['chown', '$USER', dst])
+        self._run(['chmod', '755', dst])
+
     def get_payload_for_nebraska(self,
                                  job_repo_url=None,
                                  build=None,
diff --git a/server/site_tests/autoupdate_Basic/autoupdate_Basic.py b/server/site_tests/autoupdate_Basic/autoupdate_Basic.py
index 88d3f90978..4bb4bcca3a 100644
--- a/server/site_tests/autoupdate_Basic/autoupdate_Basic.py
+++ b/server/site_tests/autoupdate_Basic/autoupdate_Basic.py
@@ -56,7 +56,7 @@ class autoupdate_Basic(update_engine_test.UpdateEngineTest):
                 public_bucket=running_at_desk)
 
         self._m2n = m2n
-        if self._m2n and not running_at_desk:
+        if self._m2n:
             if self._host.get_board().endswith("-kernelnext"):
                 raise error.TestNAError("Skipping test on kernelnext board")
 
@@ -65,14 +65,22 @@ class autoupdate_Basic(update_engine_test.UpdateEngineTest):
             logging.debug('build name is %s', build_name)
 
             # Install the matching build with quick provision.
-            autotest_devserver = dev_server.ImageServer.resolve(
-                    build_name, self._host.hostname)
-            update_url = autotest_devserver.get_update_url(build_name)
+            if running_at_desk:
+                self._copy_quick_provision_to_dut()
+                update_url = self._get_provision_url_on_public_bucket(
+                        build_name)
+            else:
+                autotest_devserver = dev_server.ImageServer.resolve(
+                        build_name, self._host.hostname)
+                update_url = autotest_devserver.get_update_url(build_name)
+
             logging.info('Installing source image with update url: %s',
                          update_url)
             provisioner.ChromiumOSProvisioner(
-                    update_url, host=self._host,
-                    is_release_bucket=True).run_provision()
+                    update_url,
+                    host=self._host,
+                    is_release_bucket=True,
+                    public_bucket=running_at_desk).run_provision()
 
         # Login to device before update
         if pin_login:
diff --git a/server/site_tests/autoupdate_Basic/control.m2n.full b/server/site_tests/autoupdate_Basic/control.m2n.full
index 3fbd75cfd0..e62298838d 100644
--- a/server/site_tests/autoupdate_Basic/control.m2n.full
+++ b/server/site_tests/autoupdate_Basic/control.m2n.full
@@ -20,23 +20,15 @@ changes don't break autoupdate.
 In the lab, a job_repo_url will be passed directly to the test. It contains
 information about the build to use and the IP address of lab cache server to
 download update payloads from. Local runs can use cache servers as well but
-the setup is fairly complicated.
-
-M2N tests require lab cache server access to provision the source stable build,
-so they are not easy to run locally. The simplest alternative is to flash a
-stable build to the DUT yourself (such as with `cros flash`) and run the
-test with the `running_at_desk` and `build` args specified.
-
-If you have lab cache server SSH set up and a valid list of IPs in your
-shadow_config.ini, you can run this test locally with a job_repo_url.
+the setup is fairly complicated. Instead, local test runs can specify the
+`running_at_desk` and `build` args to run without requiring special setup.
 
 Example usage:
 
-With cache server access and a job_repo_url:
+With lab cache server SSH access and a job_repo_url:
 test_that <DUT> autoupdate_Basic.m2n.full --args="job_repo_url='http://<IP>:<port>/static/<board>-release/RXX-XXXX.X.X/autotest/packages'"
 
-Without one, both of these args are required. You have to flash the source
-version yourself before running this:
+Otherwise, the running_at_desk and build args are required:
 test_that <DUT> autoupdate_Basic.m2n.full --board=<board> --args="running_at_desk=True build=R102-14643.0.0"
 
 """
diff --git a/server/site_tests/autoupdate_Basic/control.m2n.full.pin b/server/site_tests/autoupdate_Basic/control.m2n.full.pin
index 4a8646cdd7..5f1a795af8 100644
--- a/server/site_tests/autoupdate_Basic/control.m2n.full.pin
+++ b/server/site_tests/autoupdate_Basic/control.m2n.full.pin
@@ -19,23 +19,15 @@ changes don't break autoupdate.
 In the lab, a job_repo_url will be passed directly to the test. It contains
 information about the build to use and the IP address of lab cache server to
 download update payloads from. Local runs can use cache servers as well but
-the setup is fairly complicated.
-
-M2N tests require lab cache server access to provision the source stable build,
-so they are not easy to run locally. The simplest alternative is to flash a
-stable build to the DUT yourself (such as with `cros flash`) and run the
-test with the `running_at_desk` and `build` args specified.
-
-If you have lab cache server SSH set up and a valid list of IPs in your
-shadow_config.ini, you can run this test locally with a job_repo_url.
+the setup is fairly complicated. Instead, local test runs can specify the
+`running_at_desk` and `build` args to run without requiring special setup.
 
 Example usage:
 
-With cache server access and a job_repo_url:
+With lab cache server SSH access and a job_repo_url:
 test_that <DUT> autoupdate_Basic.m2n.full.pin --args="job_repo_url='http://<IP>:<port>/static/<board>-release/RXX-XXXX.X.X/autotest/packages'"
 
-Without one, both of these args are required. You have to flash the source
-version yourself before running this:
+Otherwise, the running_at_desk and build args are required:
 test_that <DUT> autoupdate_Basic.m2n.full.pin --board=<board> --args="running_at_desk=True build=R102-14643.0.0"
 
 """