[autotest] Support Android/Brillo suite runs.

This change does 3 key things:

1) Enables suite runs to use the production code currently
   deployed on the lab servers. This disables this requirement
   and now run_suite.py can be passed --run_prod_code. This means
   that it will use the control files on the drone and SSP has
   been disabled.
2) Refactors out the AFE interactions for applying version labels
   out of the host objects. These utilities are in a new module
   called afe_utils. The afe_utils:host_in_lab function
   determines if a host is running in an lab environment and if
   it is in the AFE.
   The introduction of afe_utils enables further cleanup of the
   cros_host object.
3) Enables provisioning for Android/Brillo devices. Added a new
   site-test and provision attribute. Run_suite will determine
   if an Android/Brillo build is being used and will set the
   correct version prefix. Then the scheduler/provision code
   will find an appropriate device and kick off the
   provision_AndroidUpdate test to install the build prior to
   test runs.

BUG=chromium:574566
TEST=On my moblab:
./run_suite.py --build=git_mnc-brillo-dev/dragonboard-userdebug/2558576 \
--board=brillo-dragonboard --suite_name=brillo-bvt --run_prod_code \
--pool=''
./site_utils/run_suite.py --board=peppy --build=peppy-release/LATEST \
--suite_name=dummy --pool=''
test_that 100.96.51.40 dummy_PassServer
DEPLOY=apache

Change-Id: I5a81e6bd989188b7c0621e2a76752a39033f9782
Reviewed-on: https://chromium-review.googlesource.com/323781
Commit-Ready: Simran Basi <sbasi@chromium.org>
Tested-by: Simran Basi <sbasi@chromium.org>
Reviewed-by: Christopher Wiley <wiley@chromium.org>
Reviewed-by: Fang Deng <fdeng@chromium.org>
Reviewed-by: Simran Basi <sbasi@chromium.org>
diff --git a/frontend/afe/site_rpc_interface.py b/frontend/afe/site_rpc_interface.py
index e0a7e0a..074e438 100644
--- a/frontend/afe/site_rpc_interface.py
+++ b/frontend/afe/site_rpc_interface.py
@@ -128,7 +128,7 @@
                      suite_args=None, wait_for_results=True, job_retry=False,
                      max_retries=None, max_runtime_mins=None, suite_min_duts=0,
                      offload_failures_only=False, builds={},
-                     test_source_build=None, **kwargs):
+                     test_source_build=None, run_prod_code=False, **kwargs):
     """
     Create a job to run a test suite on the given device with the given image.
 
@@ -170,6 +170,11 @@
                            competing with another suite that has a higher
                            priority but already got minimum machines it needs.
     @param offload_failures_only: Only enable gs_offloading for failed jobs.
+    @param run_prod_code: If True, the suite will run the test code that
+                          lives in prod aka the test code currently on the
+                          lab servers. If False, the control files and test
+                          code for this suite run will be retrieved from the
+                          build artifacts.
     @param kwargs: extra keyword args. NOT USED.
 
     @raises ControlFileNotFound: if a unique suite control file doesn't exist.
@@ -201,7 +206,16 @@
     test_source_build = Suite.get_test_source_build(
             builds, test_source_build=test_source_build)
 
-    (ds, keyvals) = _stage_build_artifacts(test_source_build)
+    suite_name = canonicalize_suite_name(name)
+    if run_prod_code:
+        ds = dev_server.ImageServer.resolve(build)
+        keyvals = {}
+        getter = control_file_getter.FileSystemGetter(
+                [_CONFIG.get_config_value('SCHEDULER',
+                                          'drone_installation_directory')])
+        control_file = getter.get_control_file_contents_by_name(suite_name)
+    else:
+        (ds, keyvals) = _stage_build_artifacts(test_source_build)
     keyvals[constants.SUITE_MIN_DUTS_KEY] = suite_min_duts
 
     if not control_file:
@@ -223,7 +237,7 @@
     # R45 falls out of stable channel.
     # Prepend build and board to the control file.
     inject_dict = {'board': board,
-                   'build': builds[provision.CROS_VERSION_PREFIX],
+                   'build': builds.get(provision.CROS_VERSION_PREFIX),
                    'builds': builds,
                    'check_hosts': check_hosts,
                    'pool': pool,
@@ -239,7 +253,8 @@
                    'max_retries': max_retries,
                    'max_runtime_mins': max_runtime_mins,
                    'offload_failures_only': offload_failures_only,
-                   'test_source_build': test_source_build
+                   'test_source_build': test_source_build,
+                   'run_prod_code': run_prod_code
                    }
 
     control_file = tools.inject_vars(inject_dict, control_file)
diff --git a/global_config.ini b/global_config.ini
index 1b324ad..c3898b4 100644
--- a/global_config.ini
+++ b/global_config.ini
@@ -378,3 +378,6 @@
 [ACTS]
 # Section for ACTS configuration.
 acts_config_folder:
+
+[ANDROID]
+image_url_pattern: %s/static/%s
diff --git a/server/afe_utils.py b/server/afe_utils.py
new file mode 100644
index 0000000..78d7595
--- /dev/null
+++ b/server/afe_utils.py
@@ -0,0 +1,84 @@
+# Copyright 2016 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Utility functions for AFE-based interactions."""
+
+import common
+from autotest_lib.server import utils
+from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
+
+AFE = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10)
+
+
+def host_in_lab(host):
+    """Check if the host is in the lab and an object the AFE knows.
+
+    This check ensures that autoserv and the host's current job is running
+    inside a fully Autotest instance, aka a lab environment. If this is the
+    case it then verifies the host is registed with the configured AFE
+    instance.
+
+    @param host: Host object to verify.
+
+    @returns The host model object.
+    """
+    if not host.job.in_lab:
+        return False
+    return AFE.get_hosts(hostname=host.hostname)
+
+
+def get_build(host):
+    """Retrieve the current build for a given hostname from the AFE.
+
+    Looks through a host's labels in the AFE to determine its build.
+
+    @param hostname: Hostname of the host whose build we want to retrieve.
+
+    @returns The current build or None if it could not find it or if there
+             were multiple build labels assigned to the host.
+    """
+    if not host_in_lab(host):
+        return None
+    return utils.get_build_from_afe(host.hostname, AFE)
+
+
+def clear_version_labels(host):
+    """Clear version labels for a given host.
+
+    @param host: Host whose version labels to clear.
+    """
+    if not host_in_lab(host):
+        return
+
+    host_list = [host.hostname]
+    labels = AFE.get_labels(
+            name__startswith=host.VERSION_PREFIX,
+            host__hostname=host.hostname)
+
+    for label in labels:
+        label.remove_hosts(hosts=host_list)
+
+
+def add_version_label(host, image_name):
+    """Add version labels to a host.
+
+    @param host: Host to add the version label for.
+    @param image_name: Name of the build version to add to the host.
+    """
+    if not host_in_lab(host):
+        return
+    label = '%s:%s' % (host.VERSION_PREFIX, image_name)
+    AFE.run('label_add_hosts', id=label, hosts=[host.hostname])
+
+
+def machine_install_and_update_labels(host, *args, **dargs):
+    """Calls machine_install and updates the version labels on a host.
+
+    @param host: Host object to run machine_install on.
+    @param *args: Args list to pass to machine_install.
+    @param **dargs: dargs dict to pass to machine_install.
+    """
+    clear_version_labels(host)
+    image_name = host.machine_install(*args, **dargs)
+    add_version_label(host, image_name)
diff --git a/server/control_segments/install b/server/control_segments/install
index ab08e3d..9650870 100644
--- a/server/control_segments/install
+++ b/server/control_segments/install
@@ -1,7 +1,10 @@
+import common
+from autotest_lib.server import afe_utils
+
 def install(machine):
     host = hosts.create_target_machine(machine, initialize=False,
                                        auto_monitor=False)
-    host.machine_install()
+    afe_utils.machine_install_and_update_labels(host)
 
 
 job.parallel_simple(install, machines, log=False)
diff --git a/server/cros/dynamic_suite/control_file_getter.py b/server/cros/dynamic_suite/control_file_getter.py
index e851020..af48419 100644
--- a/server/cros/dynamic_suite/control_file_getter.py
+++ b/server/cros/dynamic_suite/control_file_getter.py
@@ -154,14 +154,21 @@
             directory = directories.pop()
             if not os.path.exists(directory):
                 continue
-            for name in os.listdir(directory):
-                fullpath = os.path.join(directory, name)
-                if os.path.isfile(fullpath):
-                    if regexp.search(name):
-                        # if we are a control file
-                        self._files.append(fullpath)
-                elif os.path.isdir(fullpath):
-                    directories.append(fullpath)
+            try:
+                for name in os.listdir(directory):
+                    fullpath = os.path.join(directory, name)
+                    if os.path.isfile(fullpath):
+                        if regexp.search(name):
+                            # if we are a control file
+                            self._files.append(fullpath)
+                    elif os.path.isdir(fullpath):
+                        directories.append(fullpath)
+            except OSError:
+                # Some directories under results/ like the Chrome Crash
+                # Reports will cause issues when attempted to be searched.
+                logging.error('Unable to search directory %d for control '
+                              'files.', directory)
+                pass
         if not self._files:
             msg = 'No control files under ' + ','.join(self._paths)
             raise error.NoControlFileList(msg)
diff --git a/server/cros/dynamic_suite/dynamic_suite.py b/server/cros/dynamic_suite/dynamic_suite.py
index 10d1486..daf9142 100644
--- a/server/cros/dynamic_suite/dynamic_suite.py
+++ b/server/cros/dynamic_suite/dynamic_suite.py
@@ -274,7 +274,8 @@
                  bug_template={}, devserver_url=None,
                  priority=priorities.Priority.DEFAULT, predicate=None,
                  wait_for_results=True, job_retry=False, max_retries=None,
-                 offload_failures_only=False, test_source_build=None, **dargs):
+                 offload_failures_only=False, test_source_build=None,
+                 run_prod_code=False, **dargs):
         """
         Vets arguments for reimage_and_run() and populates self with supplied
         values.
@@ -350,6 +351,9 @@
                             Default to None, no max.
         @param offload_failures_only: Only enable gs_offloading for failed
                                       jobs.
+        @param run_prod_code: If true, the suite will run the test code that
+                              lives in prod aka the test code currently on the
+                              lab servers.
         @param **dargs: these arguments will be ignored.  This allows us to
                         deprecate and remove arguments in ToT while not
                         breaking branch builds.
@@ -430,6 +434,7 @@
         self.job_retry = job_retry
         self.max_retries = max_retries
         self.offload_failures_only = offload_failures_only
+        self.run_prod_code = run_prod_code
 
 
 def skip_reimage(g):
@@ -555,8 +560,9 @@
     # control_files and test_suites packages so that we can get the control
     # files we should schedule.
     try:
-        spec.devserver.stage_artifacts(spec.test_source_build,
-                                       ['control_files', 'test_suites'])
+        if not spec.run_prod_code:
+            spec.devserver.stage_artifacts(spec.test_source_build,
+                                           ['control_files', 'test_suites'])
     except dev_server.DevServerException as e:
         # If we can't get the control files, there's nothing to run.
         raise error.AsynchronousBuildFailure(e)
@@ -578,7 +584,8 @@
         priority=spec.priority, wait_for_results=spec.wait_for_results,
         job_retry=spec.job_retry, max_retries=spec.max_retries,
         offload_failures_only=spec.offload_failures_only,
-        test_source_build=spec.test_source_build)
+        test_source_build=spec.test_source_build,
+        run_prod_code=spec.run_prod_code)
 
     # Now we get to asychronously schedule tests.
     suite.schedule(spec.job.record_entry, spec.add_experimental)
diff --git a/server/cros/dynamic_suite/dynamic_suite_unittest.py b/server/cros/dynamic_suite/dynamic_suite_unittest.py
index d94c54c..df2fe6d 100755
--- a/server/cros/dynamic_suite/dynamic_suite_unittest.py
+++ b/server/cros/dynamic_suite/dynamic_suite_unittest.py
@@ -153,6 +153,7 @@
         spec.devserver.stage_artifacts(
                 spec.builds[provision.CROS_VERSION_PREFIX],
                 ['control_files', 'test_suites']).WithSideEffects(suicide)
+        spec.run_prod_code = False
 
         self.mox.ReplayAll()
 
diff --git a/server/cros/dynamic_suite/suite.py b/server/cros/dynamic_suite/suite.py
index d178c1d..50ed9f5 100644
--- a/server/cros/dynamic_suite/suite.py
+++ b/server/cros/dynamic_suite/suite.py
@@ -18,6 +18,7 @@
 from autotest_lib.client.common_lib import control_data
 from autotest_lib.client.common_lib import enum
 from autotest_lib.client.common_lib import error
+from autotest_lib.client.common_lib import global_config
 from autotest_lib.client.common_lib import priorities
 from autotest_lib.client.common_lib import site_utils
 from autotest_lib.client.common_lib import time_utils
@@ -46,6 +47,9 @@
 _FILE_BUG_SUITES = ['au', 'bvt', 'bvt-cq', 'bvt-inline', 'paygen_au_beta',
                     'paygen_au_canary', 'paygen_au_dev', 'paygen_au_stable',
                     'sanity', 'push_to_prod']
+_AUTOTEST_DIR = global_config.global_config.get_config_value(
+        'SCHEDULER', 'drone_installation_directory')
+
 
 class RetryHandler(object):
     """Maintain retry information.
@@ -514,7 +518,8 @@
 
     @staticmethod
     def create_from_predicates(predicates, builds, board, devserver,
-                               cf_getter=None, name='ad_hoc_suite', **dargs):
+                               cf_getter=None, name='ad_hoc_suite',
+                               run_prod_code=False, **dargs):
         """
         Create a Suite using a given predicate test filters.
 
@@ -534,16 +539,22 @@
         @param cf_getter: control_file_getter.ControlFileGetter. Defaults to
                           using DevServerGetter.
         @param name: name of suite. Defaults to 'ad_hoc_suite'
+        @param run_prod_code: If true, the suite will run the tests that
+                              lives in prod aka the test code currently on the
+                              lab servers.
         @param **dargs: Any other Suite constructor parameters, as described
                         in Suite.__init__ docstring.
         @return a Suite instance.
         """
         if cf_getter is None:
             build = Suite.get_test_source_build(builds, **dargs)
-            cf_getter = Suite.create_ds_getter(build, devserver)
+            if run_prod_code:
+                cf_getter = Suite.create_fs_getter(_AUTOTEST_DIR)
+            else:
+                cf_getter = Suite.create_ds_getter(build, devserver)
 
         return Suite(predicates,
-                     name, builds, board, cf_getter, **dargs)
+                     name, builds, board, cf_getter, run_prod_code, **dargs)
 
 
     @staticmethod
@@ -576,11 +587,11 @@
                      name, builds, board, cf_getter, **dargs)
 
 
-    def __init__(self, predicates, tag, builds, board, cf_getter, afe=None,
-                 tko=None, pool=None, results_dir=None, max_runtime_mins=24*60,
-                 timeout_mins=24*60, file_bugs=False,
-                 file_experimental_bugs=False, suite_job_id=None,
-                 ignore_deps=False, extra_deps=[],
+    def __init__(self, predicates, tag, builds, board, cf_getter,
+                 run_prod_code=False, afe=None, tko=None, pool=None,
+                 results_dir=None, max_runtime_mins=24*60, timeout_mins=24*60,
+                 file_bugs=False, file_experimental_bugs=False,
+                 suite_job_id=None, ignore_deps=False, extra_deps=[],
                  priority=priorities.Priority.DEFAULT, forgiving_parser=True,
                  wait_for_results=True, job_retry=False,
                  max_retries=sys.maxint, offload_failures_only=False,
@@ -600,6 +611,9 @@
         @param tko: an instance of TKO as defined in server/frontend.py.
         @param pool: Specify the pool of machines to use for scheduling
                 purposes.
+        @param run_prod_code: If true, the suite will run the test code that
+                              lives in prod aka the test code currently on the
+                              lab servers.
         @param results_dir: The directory where the job can write results to.
                             This must be set if you want job_id of sub-jobs
                             list in the job keyvals.
@@ -651,7 +665,8 @@
         self._jobs_to_tests = {}
         self._tests = Suite.find_and_parse_tests(self._cf_getter,
                         self._predicate, self._tag, add_experimental=True,
-                        forgiving_parser=forgiving_parser)
+                        forgiving_parser=forgiving_parser,
+                        run_prod_code=run_prod_code)
 
         self._max_runtime_mins = max_runtime_mins
         self._timeout_mins = timeout_mins
@@ -1076,7 +1091,7 @@
 
     @staticmethod
     def find_all_tests(cf_getter, suite_name='', add_experimental=False,
-                       forgiving_parser=True):
+                       forgiving_parser=True, run_prod_code=False):
         """
         Function to scan through all tests and find all tests.
 
@@ -1100,6 +1115,10 @@
                                  files. Note that this can raise an exception
                                  for syntax errors in unrelated files, because
                                  we parse them before applying the predicate.
+        @param run_prod_code: If true, the suite will run the test code that
+                              lives in prod aka the test code currently on the
+                              lab servers by disabling SSP for the discovered
+                              tests.
 
         @raises ControlVariableException: If forgiving_parser is False and there
                                           is a syntax error in a control file.
@@ -1122,6 +1141,8 @@
                     continue
                 found_test.text = text
                 found_test.path = file
+                if run_prod_code:
+                    found_test.require_ssp = False
                 tests[file] = found_test
             except control_data.ControlVariableException, e:
                 if not forgiving_parser:
@@ -1135,7 +1156,8 @@
 
     @staticmethod
     def find_and_parse_tests(cf_getter, predicate, suite_name='',
-                             add_experimental=False, forgiving_parser=True):
+                             add_experimental=False, forgiving_parser=True,
+                             run_prod_code=False):
         """
         Function to scan through all tests and find eligible tests.
 
@@ -1156,6 +1178,10 @@
                                  files. Note that this can raise an exception
                                  for syntax errors in unrelated files, because
                                  we parse them before applying the predicate.
+        @param run_prod_code: If true, the suite will run the test code that
+                              lives in prod aka the test code currently on the
+                              lab servers by disabling SSP for the discovered
+                              tests.
 
         @raises ControlVariableException: If forgiving_parser is False and there
                                           is a syntax error in a control file.
@@ -1165,7 +1191,8 @@
                 on the TIME setting in control file, slowest test comes first.
         """
         tests = Suite.find_all_tests(cf_getter, suite_name, add_experimental,
-                                     forgiving_parser)
+                                     forgiving_parser,
+                                     run_prod_code=run_prod_code)
         logging.debug('Parsed %s control files.', len(tests))
         tests = [test for test in tests.itervalues() if predicate(test)]
         tests.sort(key=lambda t:
diff --git a/server/cros/dynamic_suite/suite_unittest.py b/server/cros/dynamic_suite/suite_unittest.py
index 41aa7b0..f8dd92f 100755
--- a/server/cros/dynamic_suite/suite_unittest.py
+++ b/server/cros/dynamic_suite/suite_unittest.py
@@ -237,7 +237,8 @@
             mox.IgnoreArg(),
             mox.IgnoreArg(),
             add_experimental=True,
-            forgiving_parser=True).AndReturn(self.files.values())
+            forgiving_parser=True,
+            run_prod_code=False).AndReturn(self.files.values())
 
 
     def expect_job_scheduling(self, recorder, add_experimental,
diff --git a/server/cros/provision.py b/server/cros/provision.py
index 96a0081..c12f48f 100644
--- a/server/cros/provision.py
+++ b/server/cros/provision.py
@@ -11,6 +11,7 @@
 
 ### Constants for label prefixes
 CROS_VERSION_PREFIX = 'cros-version'
+ANDROID_BUILD_VERSION_PREFIX = 'ab-version'
 FW_RW_VERSION_PREFIX = 'fwrw-version'
 FW_RO_VERSION_PREFIX = 'fwro-version'
 
@@ -171,6 +172,8 @@
         FW_RW_VERSION_PREFIX: actionables.TestActionable(
                 'provision_FirmwareUpdate',
                 extra_kwargs={'rw_only': True}),
+        ANDROID_BUILD_VERSION_PREFIX : actionables.TestActionable(
+                'provision_AndroidUpdate'),
     }
 
     name = 'provision'
diff --git a/server/cros/telemetry_runner.py b/server/cros/telemetry_runner.py
index 15eeb6d..30c44a5 100644
--- a/server/cros/telemetry_runner.py
+++ b/server/cros/telemetry_runner.py
@@ -12,6 +12,7 @@
 
 from autotest_lib.client.common_lib import error, utils
 from autotest_lib.client.common_lib.cros import dev_server
+from autotest_lib.server import afe_utils
 
 
 TELEMETRY_RUN_BENCHMARKS_SCRIPT = 'tools/perf/run_benchmark'
@@ -271,7 +272,7 @@
         logging.debug('Setting up telemetry for devserver testing')
         logging.debug('Grabbing build from AFE.')
 
-        build = self._host.get_build()
+        build = afe_utils.get_build(self._host)
         if not build:
             logging.error('Unable to locate build label for host: %s.',
                           self._host.hostname)
diff --git a/server/hosts/abstract_ssh.py b/server/hosts/abstract_ssh.py
index d8f51b3..4854d9e 100644
--- a/server/hosts/abstract_ssh.py
+++ b/server/hosts/abstract_ssh.py
@@ -22,6 +22,7 @@
     almost all of the abstract Host methods, except for the core
     Host.run method.
     """
+    VERSION_PREFIX = ''
 
     def _initialize(self, hostname, user="root", port=22, password="",
                     is_client_install_supported=True, host_attributes={},
@@ -855,4 +856,4 @@
 
         @return A string describing the OS type.
         """
-        raise NotImplementedError
+        raise NotImplementedError
\ No newline at end of file
diff --git a/server/hosts/adb_host.py b/server/hosts/adb_host.py
index 1998452..4df58af 100644
--- a/server/hosts/adb_host.py
+++ b/server/hosts/adb_host.py
@@ -17,6 +17,7 @@
 from autotest_lib.server import autoserv_parser
 from autotest_lib.server import constants as server_constants
 from autotest_lib.server import utils
+from autotest_lib.server.cros import provision
 from autotest_lib.server.cros.dynamic_suite import constants
 from autotest_lib.server.hosts import abstract_ssh
 from autotest_lib.server.hosts import teststation_host
@@ -104,6 +105,7 @@
 class ADBHost(abstract_ssh.AbstractSSHHost):
     """This class represents a host running an ADB server."""
 
+    VERSION_PREFIX = provision.ANDROID_BUILD_VERSION_PREFIX
     _LABEL_FUNCTIONS = []
     _DETECTABLE_LABELS = []
     label_decorator = functools.partial(utils.add_label_detector,
@@ -1104,10 +1106,6 @@
         # folder used to store image files after the provision is completed.
         delete_build_folder = bool(not build_local_path)
 
-        if not build_url and self._parser.options.image:
-            build_url, _ = self.stage_build_for_install(
-                    self._parser.options.image)
-
         try:
             # Download image files needed for provision to a local directory.
             if not build_local_path:
@@ -1169,10 +1167,6 @@
         # folder used to store image files after the provision is completed.
         delete_build_folder = bool(not build_local_path)
 
-        if not build_url and self._parser.options.image:
-            build_url, _ = self.stage_build_for_install(
-                    self._parser.options.image)
-
         try:
             # Download image files needed for provision to a local directory.
             if not build_local_path:
@@ -1213,7 +1207,12 @@
         @param wipe: If true, userdata will be wiped before flashing.
         @param flash_all: If True, all img files found in img_path will be
                 flashed. Otherwise, only boot and system are flashed.
+
+        @returns Name of the image installed.
         """
+        if not build_url and self._parser.options.image:
+            build_url, _ = self.stage_build_for_install(
+                    self._parser.options.image)
         if self.get_os_type() == OS_TYPE_ANDROID:
             self.install_android(
                     build_url=build_url, build_local_path=build_local_path,
@@ -1225,6 +1224,7 @@
             raise error.InstallError(
                     'Installation of os type %s is not supported.' %
                     self.get_os_type())
+        return build_url.split('static/')[-1]
 
 
     def list_files_glob(self, path_glob):
diff --git a/server/hosts/cros_host.py b/server/hosts/cros_host.py
index c33d7b9..416f90d 100644
--- a/server/hosts/cros_host.py
+++ b/server/hosts/cros_host.py
@@ -24,6 +24,7 @@
 from autotest_lib.client.cros.audio import cras_utils
 from autotest_lib.client.cros.input_playback import input_playback
 from autotest_lib.client.cros.video import constants as video_test_constants
+from autotest_lib.server import afe_utils
 from autotest_lib.server import autoserv_parser
 from autotest_lib.server import autotest
 from autotest_lib.server import constants
@@ -57,6 +58,8 @@
 class CrosHost(abstract_ssh.AbstractSSHHost):
     """Chromium OS specific subclass of Host."""
 
+    VERSION_PREFIX = provision.CROS_VERSION_PREFIX
+
     _parser = autoserv_parser.autoserv_parser
     _AFE = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10)
 
@@ -378,14 +381,6 @@
         return stable_version
 
 
-    def _host_in_AFE(self):
-        """Check if the host is an object the AFE knows.
-
-        @returns the host object.
-        """
-        return self._AFE.get_hosts(hostname=self.hostname)
-
-
     def lookup_job_repo_url(self):
         """Looks up the job_repo_url for the host.
 
@@ -400,19 +395,10 @@
             return None
 
 
-    def clear_cros_version_labels_and_job_repo_url(self):
-        """Clear cros_version labels and host attribute job_repo_url."""
-        if not self._host_in_AFE():
+    def clear_job_repo_url(self):
+        """Clear host attribute job_repo_url."""
+        if not afe_utils.host_in_lab(self):
             return
-
-        host_list = [self.hostname]
-        labels = self._AFE.get_labels(
-                name__startswith=ds_constants.VERSION_PREFIX,
-                host__hostname=self.hostname)
-
-        for label in labels:
-            label.remove_hosts(hosts=host_list)
-
         self.update_job_repo_url(None, None)
 
 
@@ -435,21 +421,18 @@
                                       'host %s' % (repo_url, self.hostname))
 
 
-    def add_cros_version_labels_and_job_repo_url(self, image_name):
+    def add_job_repo_url(self, image_name):
         """Add cros_version labels and host attribute job_repo_url.
 
         @param image_name: The name of the image e.g.
                 lumpy-release/R27-3837.0.0
 
         """
-        if not self._host_in_AFE():
+        if not afe_utils.host_in_lab(self):
             return
 
-        cros_label = '%s%s' % (ds_constants.VERSION_PREFIX, image_name)
         devserver_url = dev_server.ImageServer.resolve(image_name,
                                                        self.hostname).url()
-
-        self._AFE.run('label_add_hosts', id=cros_label, hosts=[self.hostname])
         self.update_job_repo_url(devserver_url, image_name)
 
 
@@ -723,9 +706,8 @@
         failed to do a stateful update, full update, including kernel update,
         will be applied to the DUT.
 
-        Once a host enters machine_install its cros_version label will be
-        removed as well as its host attribute job_repo_url (used for
-        package install).
+        Once a host enters machine_install its host attribute job_repo_url
+        (used for package install) will be removed and then updated.
 
         @param update_url: The url to use for the update
                 pattern: http://$devserver:###/update/$build
@@ -743,6 +725,7 @@
                 first when the dut is already installed with the same version.
         @raises autoupdater.ChromiumOSError
 
+        @returns Name of the image installed.
         """
         devserver = None
         if repair:
@@ -768,7 +751,7 @@
         logging.debug('Update URL is %s', update_url)
 
         # Remove cros-version and job_repo_url host attribute from host.
-        self.clear_cros_version_labels_and_job_repo_url()
+        self.clear_job_repo_url()
 
         # Create a file to indicate if provision fails. The file will be removed
         # by stateful update or full install.
@@ -838,8 +821,9 @@
             self.reboot(timeout=self.REBOOT_TIMEOUT, wait=True)
 
         self._post_update_processing(updater, inactive_kernel)
-        self.add_cros_version_labels_and_job_repo_url(
-                autoupdater.url_to_image_name(update_url))
+        image_name = autoupdater.url_to_image_name(update_url)
+        self.add_job_repo_url(image_name)
+        return image_name
 
 
     def _clear_fw_version_labels(self, rw_only):
@@ -966,17 +950,6 @@
         return server_utils.get_board_from_afe(self.hostname, self._AFE)
 
 
-    def get_build(self):
-        """Retrieve the current build for this Host from the AFE.
-
-        Looks through this host's labels in the AFE to determine its build.
-
-        @returns The current build or None if it could not find it or if there
-                 were multiple build labels assigned to this host.
-        """
-        return server_utils.get_build_from_afe(self.hostname, self._AFE)
-
-
     def _install_repair(self):
         """Attempt to repair this host using the update-engine.
 
@@ -993,7 +966,7 @@
             raise error.AutoservRepairMethodNA('DUT unreachable for install.')
         logging.info('Attempting to reimage machine to repair image.')
         try:
-            self.machine_install(repair=True)
+            afe_utils.machine_install_and_update_labels(self, repair=True)
         except autoupdater.ChromiumOSError as e:
             logging.exception(e)
             logging.info('Repair via install failed.')
diff --git a/server/hosts/sonic_host.py b/server/hosts/sonic_host.py
index 9b247b3..2364a1f 100644
--- a/server/hosts/sonic_host.py
+++ b/server/hosts/sonic_host.py
@@ -267,7 +267,10 @@
 
 
     def machine_install(self, update_url):
-        """Installs a build on the Sonic device."""
+        """Installs a build on the Sonic device.
+
+        @returns String of the current build number.
+        """
         old_build_number = self.get_build_number()
         self._remount_root(permissions='rw')
         self._setup_coredump_dirs()
@@ -284,3 +287,5 @@
             raise error.AutoservRunError('Build number did not change on: '
                                          '%s after update with %s' %
                                          (self.hostname, update_url()))
+
+        return str(new_build_number)
diff --git a/server/hosts/testbed.py b/server/hosts/testbed.py
index 1891be8..e2d5987 100644
--- a/server/hosts/testbed.py
+++ b/server/hosts/testbed.py
@@ -30,6 +30,7 @@
     """This class represents a collection of connected teststations and duts."""
 
     _parser = autoserv_parser.autoserv_parser
+    VERSION_PREFIX = 'testbed-version'
 
     def __init__(self, hostname='localhost', host_attributes={},
                  adb_serials=None, **dargs):
@@ -270,7 +271,10 @@
 
 
     def machine_install(self):
-        """Install the DUT."""
+        """Install the DUT.
+
+        @returns The name of the image installed.
+        """
         if not self._parser.options.image:
             raise error.InstallError('No image string is provided to test bed.')
         images = self._parse_image(self._parser.options.image)
@@ -287,3 +291,4 @@
         thread_pool = pool.ThreadPool(_POOL_SIZE)
         thread_pool.map(self._install_device, arguments)
         thread_pool.close()
+        return self._parser.options.image
diff --git a/server/site_tests/platform_CryptohomeMigrateChapsToken/platform_CryptohomeMigrateChapsToken.py b/server/site_tests/platform_CryptohomeMigrateChapsToken/platform_CryptohomeMigrateChapsToken.py
index c2ff849..7e47c3a 100644
--- a/server/site_tests/platform_CryptohomeMigrateChapsToken/platform_CryptohomeMigrateChapsToken.py
+++ b/server/site_tests/platform_CryptohomeMigrateChapsToken/platform_CryptohomeMigrateChapsToken.py
@@ -2,7 +2,9 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-from autotest_lib.server import autotest, test
+from autotest_lib.server import afe_utils
+from autotest_lib.server import autotest
+from autotest_lib.server import test
 
 class platform_CryptohomeMigrateChapsToken(test.test):
     """ This test checks to see if Chaps generated keys are
@@ -16,7 +18,7 @@
     def run_once(self, host, baseline_version=None):
         # Save the build on the DUT, because we want to provision it after
         # the test.
-        final_version = host.get_build()
+        final_version = afe_utils.get_build(host)
         if baseline_version:
             version = baseline_version
         else:
diff --git a/server/site_tests/provision_AndroidUpdate/control b/server/site_tests/provision_AndroidUpdate/control
new file mode 100644
index 0000000..93d5d53
--- /dev/null
+++ b/server/site_tests/provision_AndroidUpdate/control
@@ -0,0 +1,39 @@
+# Copyright 2016 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+AUTHOR = "sbasi"
+NAME = "provision_AndroidUpdate"
+PURPOSE = "Provision an android-based system to the correct OS version."
+TIME = "MEDIUM"
+TEST_CATEGORY = "System"
+TEST_CLASS = "provision"
+TEST_TYPE = "Server"
+
+DOC = """
+This is a test used by the provision control segment in autoserv to set the
+ab-version label of a host to the desired setting and reimage the host to a
+specific version.
+"""
+
+
+from autotest_lib.client.common_lib import error, utils
+from autotest_lib.client.cros import constants
+
+
+# Autoserv may inject a local variable called value to supply the desired
+# version. If it does not exist, check if it was supplied as a test arg.
+if not locals().get('value'):
+    args = utils.args_to_dict(args)
+    if not args.get('value'):
+        raise error.TestError("No provision value!")
+    value = args['value']
+
+
+def run(machine):
+    host = hosts.create_host(machine, initialize=False)
+    job.run_test('provision_AndroidBuildUpdate', host=host, value=value)
+
+
+job.parallel_simple(run, machines)
diff --git a/server/site_tests/provision_AndroidUpdate/provision_AndroidUpdate.py b/server/site_tests/provision_AndroidUpdate/provision_AndroidUpdate.py
new file mode 100644
index 0000000..47dbc9c
--- /dev/null
+++ b/server/site_tests/provision_AndroidUpdate/provision_AndroidUpdate.py
@@ -0,0 +1,75 @@
+# Copyright 2016 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import logging
+
+from autotest_lib.client.common_lib import error
+from autotest_lib.client.common_lib import global_config
+from autotest_lib.server import afe_utils
+from autotest_lib.server import test
+
+
+_CONFIG = global_config.global_config
+# pylint: disable-msg=E1120
+_IMAGE_URL_PATTERN = _CONFIG.get_config_value(
+        'ANDROID', 'image_url_pattern', type=str)
+
+
+class provision_AndroidUpdate(test.test):
+    """A test that can provision a machine to the correct Android version."""
+    version = 1
+
+    def initialize(self, host, value, force=False, is_test_na=False):
+        """Initialize.
+
+        @param host: The host object to update to |value|.
+        @param value: String of the image we want to install on the host.
+        @param force: not used by initialize.
+        @param is_test_na: boolean, if True, will simply skip the test
+                           and emit TestNAError. The control file
+                           determines whether the test should be skipped
+                           and passes the decision via this argument. Note
+                           we can't raise TestNAError in control file as it won't
+                           be caught and handled properly.
+        """
+        if is_test_na:
+            raise error.TestNAError('Provisioning not applicable.')
+        # We check value in initialize so that it fails faster.
+        if not value:
+            raise error.TestFail('No build version specified.')
+
+
+    def run_once(self, host, value, force=False):
+        """The method called by the control file to start the test.
+
+        @param host: The host object to update to |value|.
+        @param value: The host object to provision with a build corresponding
+                      to |value|.
+        @param force: True iff we should re-provision the machine regardless of
+                      the current image version.  If False and the image
+                      version matches our expected image version, no
+                      provisioning will be done.
+
+        """
+        logging.debug('Start provisioning %s to %s', host, value)
+
+        # If the host is already on the correct build, we have nothing to do.
+        if not force and afe_utils.get_build(host) == value:
+            # We can't raise a TestNA, as would make sense, as that makes
+            # job.run_test return False as if the job failed.  However, it'd
+            # still be nice to get this into the status.log, so we manually
+            # emit an INFO line instead.
+            self.job.record('INFO', None, None,
+                            'Host already running %s' % value)
+            return
+
+        url, _ = host.stage_build_for_install(value)
+
+        logging.debug('Installing image from: %s', url)
+        try:
+            afe_utils.machine_install_and_update_labels(host, build_url=url)
+        except error.InstallError as e:
+            logging.error(e)
+            raise error.TestFail(str(e))
+        logging.debug('Finished provisioning %s to %s', host, value)
diff --git a/server/site_tests/provision_AutoUpdate/control.double b/server/site_tests/provision_AutoUpdate/control.double
index 382b407..77b6aba 100644
--- a/server/site_tests/provision_AutoUpdate/control.double
+++ b/server/site_tests/provision_AutoUpdate/control.double
@@ -40,6 +40,7 @@
 
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.cros import constants
+from autotest_lib.server import afe_utils
 from autotest_lib.server import utils as server_utils
 
 
@@ -59,7 +60,7 @@
         # Save preserved log after autoupdate is completed.
         job.sysinfo.add_logdir(constants.AUTOUPDATE_PRESERVE_LOG)
         host = hosts.create_host(machine)
-        value = host.get_build()
+        value = afe_utils.get_build(host)
     job.run_test('provision_AutoUpdate', host=host, value=value, force=True,
                  disable_sysinfo=False, tag='double', is_test_na=is_test_na)
 
diff --git a/server/site_tests/provision_AutoUpdate/provision_AutoUpdate.py b/server/site_tests/provision_AutoUpdate/provision_AutoUpdate.py
index 15b4a68..d89a67a 100644
--- a/server/site_tests/provision_AutoUpdate/provision_AutoUpdate.py
+++ b/server/site_tests/provision_AutoUpdate/provision_AutoUpdate.py
@@ -10,6 +10,7 @@
 from autotest_lib.client.common_lib import utils
 from autotest_lib.client.common_lib.cros import dev_server
 from autotest_lib.client.common_lib.cros.graphite import autotest_stats
+from autotest_lib.server import afe_utils
 from autotest_lib.server import test
 
 
@@ -79,10 +80,12 @@
         """The method called by the control file to start the test.
 
         @param host: The host object to update to |value|.
-        @param value: The build type and version to install on the host.
-        @param force: If False, will only provision the host if it is not
-                      already running the build. If True, force the
-                      provisioning regardless, and force a full-reimage.
+        @param value: The host object to provision with a build corresponding
+                      to |value|.
+        @param force: True iff we should re-provision the machine regardless of
+                      the current image version.  If False and the image
+                      version matches our expected image version, no
+                      provisioning will be done.
 
         """
         logging.debug('Start provisioning %s to %s', host, value)
@@ -94,7 +97,7 @@
         # We could just not pass |force_update=True| to |machine_install|,
         # but I'd like the semantics that a provision test 'returns' TestNA
         # if the machine is already properly provisioned.
-        if not force and host.get_build() == value:
+        if not force and afe_utils.get_build(host) == value:
             # We can't raise a TestNA, as would make sense, as that makes
             # job.run_test return False as if the job failed.  However, it'd
             # still be nice to get this into the status.log, so we manually
@@ -127,8 +130,10 @@
 
         logging.debug('Installing image')
         try:
-            host.machine_install(force_update=True, update_url=url,
-                                 force_full_update=force)
+            afe_utils.machine_install_and_update_labels(host,
+                                                        force_update=True,
+                                                        update_url=url,
+                                                        force_full_update=force)
         except error.InstallError as e:
             logging.error(e)
             raise error.TestFail(str(e))
diff --git a/server/tests/reinstall/reinstall.py b/server/tests/reinstall/reinstall.py
index 1060ab4..6276d07 100644
--- a/server/tests/reinstall/reinstall.py
+++ b/server/tests/reinstall/reinstall.py
@@ -1,12 +1,13 @@
-import time
-from autotest_lib.server import test
+import common
 from autotest_lib.client.common_lib import error
+from autotest_lib.server import afe_utils
+from autotest_lib.server import test
 
 class reinstall(test.test):
     version = 1
 
     def execute(self, host):
         try:
-            host.machine_install()
+            afe_utils.machine_install_and_update_labels(host)
         except Exception, e:
             raise error.TestFail(str(e))
diff --git a/site_utils/run_suite.py b/site_utils/run_suite.py
index c9b6d5a..658a566 100755
--- a/site_utils/run_suite.py
+++ b/site_utils/run_suite.py
@@ -80,6 +80,7 @@
             RETURN_CODES.SUITE_TIMEOUT: 2,
             RETURN_CODES.INFRA_FAILURE: 3,
             RETURN_CODES.ERROR: 4}
+ANDROID_BUILD_REGEX = r'.+/.+/[0-9]+'
 
 
 def get_worse_code(code1, code2):
@@ -187,6 +188,10 @@
     parser.add_option('--json_dump', dest='json_dump', action='store_true',
                       default=False,
                       help='Dump the output of run_suite to stdout.')
+    parser.add_option('--run_prod_code', dest='run_prod_code',
+                      action='store_true', default=False,
+                      help='Run the test code that lives in prod aka the test '
+                           'code currently on the lab servers.')
     options, args = parser.parse_args()
     return parser, options, args
 
@@ -1405,7 +1410,10 @@
     """
     builds = {}
     if options.build:
-        builds[provision.CROS_VERSION_PREFIX] = options.build
+        if re.match(ANDROID_BUILD_REGEX, options.build):
+            builds[provision.ANDROID_BUILD_VERSION_PREFIX] = options.build
+        else:
+            builds[provision.CROS_VERSION_PREFIX] = options.build
     if options.firmware_rw_build:
         builds[provision.FW_RW_VERSION_PREFIX] = options.firmware_rw_build
     if options.firmware_ro_build:
@@ -1437,7 +1445,8 @@
                    max_runtime_mins=options.max_runtime_mins,
                    job_retry=retry, max_retries=options.max_retries,
                    suite_min_duts=options.suite_min_duts,
-                   offload_failures_only=offload_failures_only)
+                   offload_failures_only=offload_failures_only,
+                   run_prod_code=options.run_prod_code)
 
 
 def main_without_exception_handling(options):
diff --git a/site_utils/test_push.py b/site_utils/test_push.py
index 052849d..2776f6f 100755
--- a/site_utils/test_push.py
+++ b/site_utils/test_push.py
@@ -37,6 +37,7 @@
     # Unittest may not have Django database configured and will fail to import.
     pass
 from autotest_lib.client.common_lib import global_config
+from autotest_lib.server import afe_utils
 from autotest_lib.server import site_utils
 from autotest_lib.server.cros import provision
 from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
@@ -309,7 +310,7 @@
     hostnames = set([hqe.host.hostname for hqe in hqes])
     for hostname in hostnames:
         host = cros_host.CrosHost(hostname)
-        found_build = host.get_build()
+        found_build = afe_utils.get_build(hostname)
         if found_build != build:
             raise TestPushException('DUT is not imaged properly. Host %s has '
                                     'build %s, while build %s is expected.' %