provision_au: Changing test to call chromite AU instead of autotest AU

This CL updates the provision_AutoUpdate testcase to call
chromite.lib.auto_updater - unless that fails, it reverts back to quick
provision so as to not cause an outage.

BUG=chromium:1001312
TEST=Ran inside chroot: provision_AutoUpdate with versions R75, R81.
TEST=Ran inside chroot: provision_AutoUpdate where chromite was at an
older commit.

Change-Id: I8445ff0d044e4f1f861f3906237206a79a53fb48
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/autotest/+/2015460
Commit-Queue: Sanika Kulkarni <sanikak@chromium.org>
Tested-by: Sanika Kulkarni <sanikak@chromium.org>
Reviewed-by: Allen Li <ayatane@chromium.org>
Auto-Submit: Sanika Kulkarni <sanikak@chromium.org>
diff --git a/server/afe_utils.py b/server/afe_utils.py
index e81f2d4..1c3f80a 100644
--- a/server/afe_utils.py
+++ b/server/afe_utils.py
@@ -12,12 +12,19 @@
 import common
 import logging
 import traceback
+import urlparse
+
 from autotest_lib.client.common_lib import global_config
 from autotest_lib.server.cros import autoupdater
 from autotest_lib.server.cros import provision
 from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
 from autotest_lib.site_utils import stable_version_classify as sv
 from autotest_lib.server import site_utils as server_utils
+from autotest_lib.server.cros.dynamic_suite import constants as ds_constants
+from autotest_lib.server.cros.dynamic_suite import tools
+
+from chromite.lib import auto_updater
+from chromite.lib import remote_access
 
 
 AFE = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10)
@@ -178,7 +185,7 @@
 
 def machine_install_and_update_labels(host, update_url,
                                       use_quick_provision=False,
-                                      with_cheets=False):
+                                      with_cheets=False, staging_server=None):
     """Install a build and update the version labels on a host.
 
     @param host: Host object where the build is to be installed.
@@ -187,13 +194,34 @@
         quick-provision for the update.
     @param with_cheets: If true, installation is for a specific, custom
         version of Android for a target running ARC.
+    @param staging_server: Sever where images have been staged. Typically,
+        an instance of dev_server.ImageServer.
     """
     clean_provision_labels(host)
-    updater = autoupdater.ChromiumOSUpdater(
+    # TODO(crbug.com/1049346): The try-except block exists to catch failures in
+    # chromite auto_updater that may occur due to autotest/chromite version
+    # mismatch. This should be removed once that bug is resolved.
+    try:
+        # Get image_name in the format <board>-release/Rxx-12345.0.0 from the
+        # update_url.
+        image_name = '/'.join(urlparse.urlparse(update_url).path.split('/')[-2:])
+        with remote_access.ChromiumOSDeviceHandler(host.ip) as device:
+            updater = auto_updater.ChromiumOSUpdater(
+                device, build_name=None, payload_dir=image_name,
+                staging_server=staging_server.url())
+            updater.CheckPayloads()
+            updater.PreparePayloadPropsFile()
+            updater.RunUpdate()
+        repo_url = tools.get_package_url(staging_server.url(), image_name)
+        host_attributes = {ds_constants.JOB_REPO_URL: repo_url}
+    except Exception as e:
+        logging.warning(
+            "Chromite auto_updater has failed with the exception: %s", e)
+        logging.debug("Attempting to provision with quick provision.")
+        updater = autoupdater.ChromiumOSUpdater(
             update_url, host=host, use_quick_provision=use_quick_provision)
-    image_name, host_attributes = updater.run_update()
+        image_name, host_attributes = updater.run_update()
     if with_cheets:
         image_name += provision.CHEETS_SUFFIX
-
-    add_provision_labels(
-            host, host.VERSION_PREFIX, image_name, host_attributes)
+    add_provision_labels(host, host.VERSION_PREFIX, image_name,
+                         host_attributes)
diff --git a/server/site_tests/provision_AutoUpdate/provision_AutoUpdate.py b/server/site_tests/provision_AutoUpdate/provision_AutoUpdate.py
index eb11587..f0a46f5 100644
--- a/server/site_tests/provision_AutoUpdate/provision_AutoUpdate.py
+++ b/server/site_tests/provision_AutoUpdate/provision_AutoUpdate.py
@@ -176,7 +176,8 @@
         failure = None
         try:
             afe_utils.machine_install_and_update_labels(
-                    host, url, not force_update_engine, with_cheets)
+                    host, url, not force_update_engine, with_cheets,
+                    staging_server=ds)
         except BaseException as e:
             failure = e
             raise