autotest: Add suite support for testbeds.

Adds basic suite support for Android testbeds. Note this does not
support SSP and only uses the Autotest code that is currently
pushed to prod.

* Added a new provision test (provision_TestbedUpdate).
* Added a basic suite to test this workflow.
* Updated the parsers to allow for multiple boards.
* Had to adjust the provision control_segment to allow commas
  in the key-value labels.

BUG=chromium:628415,chromium:628036
TEST=./site_utils/run_suite.py --board=dragonboard-2 \
--build=git_mnc-brillo-dev/dragonboard-eng/3014741,\
git_mnc-brillo-dev/dragonboard-eng/3014741 \
--suite_name=dummy_testbed --pool='' --run_prod_code

Change-Id: Ica04911974d69498877b9cedb9cefc4a5d0bcdd5
Reviewed-on: https://chromium-review.googlesource.com/360663
Commit-Ready: Simran Basi <sbasi@chromium.org>
Tested-by: Simran Basi <sbasi@chromium.org>
Reviewed-by: Kevin Cheng <kevcheng@chromium.org>
Reviewed-by: Dan Shi <dshi@google.com>
diff --git a/server/control_segments/provision b/server/control_segments/provision
index 2e8c65d..0276957 100644
--- a/server/control_segments/provision
+++ b/server/control_segments/provision
@@ -2,16 +2,31 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import re
 
 from autotest_lib.client.cros import constants
 from autotest_lib.server import utils
 from autotest_lib.server.cros import provision
 
+LABEL_REGEX = r',.*:'
 
 # job_labels should be a string like "name:setting,name:setting"
+# However setting might also contain ',' therefore we need more advanced logic
+# than split.
 # non-provisionable labels are currently skipped, so they're safe to pass in.
 job_labels = locals().get('job_labels') or ','.join(args)
-labels_list = [label.strip() for label in job_labels.split(',') if label]
+labels_list = []
+while job_labels:
+    # Split based off of a comma followed by colon regex.
+    split = re.split(LABEL_REGEX, job_labels)
+    # First value found is a proper key value pair.
+    labels_list.append(split[0].strip())
+    # Remove this key value pair.
+    job_labels = job_labels[len(split[0]):]
+    # If a comma remains at the start of the remaining labels, remove it.
+    # This should happen on every loop except the last one.
+    if job_labels.startswith(','):
+        job_labels = job_labels.lstrip(',')
 
 
 def provision_machine(machine):
@@ -19,7 +34,7 @@
     Run the appropriate provisioning tests to make the machine's labels match
     those given in job_labels.
     """
-    host = hosts.create_host(machine, try_lab_servo=True)
+    host = hosts.create_target_machine(machine, try_lab_servo=True)
 
     job.record('START', None, 'provision')
     try:
diff --git a/server/cros/provision.py b/server/cros/provision.py
index 6080cf5..9eb414a 100644
--- a/server/cros/provision.py
+++ b/server/cros/provision.py
@@ -212,6 +212,8 @@
                               'tag': 'rw_only'}),
         ANDROID_BUILD_VERSION_PREFIX : actionables.TestActionable(
                 'provision_AndroidUpdate'),
+        TESTBED_BUILD_VERSION_PREFIX : actionables.TestActionable(
+                'provision_TestbedUpdate'),
     }
 
     name = 'provision'
diff --git a/server/hosts/testbed.py b/server/hosts/testbed.py
index a35c17b..c6076c0 100644
--- a/server/hosts/testbed.py
+++ b/server/hosts/testbed.py
@@ -307,9 +307,12 @@
         return build_url, build_local_path, teststation
 
 
-    def machine_install(self):
+    def machine_install(self, image=None):
         """Install the DUT.
 
+        @param image: Image we want to install on this testbed, e.g.,
+                      `branch1/shamu-eng/1001,branch2/shamu-eng/1002`
+
         @returns A tuple of (the name of the image installed, None), where None
                 is a placeholder for update_url. Testbed does not have a single
                 update_url, thus it's set to None.
@@ -326,9 +329,10 @@
                 {'job_repo_url_XZ001': 'http://10.1.1.3/branch1/shamu-eng/1001',
                  'job_repo_url_XZ002': 'http://10.1.1.3/branch2/shamu-eng/1002'}
         """
-        if not self._parser.options.image:
+        image = image or self._parser.options.image
+        if not image:
             raise error.InstallError('No image string is provided to test bed.')
-        images = self._parse_image(self._parser.options.image)
+        images = self._parse_image(image)
         host_attributes = {}
 
         # Change logging formatter to include thread name. This is to help logs
@@ -365,7 +369,7 @@
                               teststation.hostname, build_local_path)
                 teststation.run('rm -rf %s' % build_local_path)
 
-        return self._parser.options.image, host_attributes
+        return image, host_attributes
 
 
     def get_attributes_to_clear_before_provision(self):
diff --git a/server/site_tests/provision_TestbedUpdate/control b/server/site_tests/provision_TestbedUpdate/control
new file mode 100644
index 0000000..497eda9
--- /dev/null
+++ b/server/site_tests/provision_TestbedUpdate/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_TestbedUpdate"
+PURPOSE = "Provision the multiple DUTs in a Testbed setup."
+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
+testbed-version label of a testbed to the desired setting and reimage the
+testbed 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):
+    testbed = hosts.create_target_machine(machine)
+    job.run_test('provision_TestbedUpdate', host=testbed, value=value)
+
+
+job.parallel_simple(run, machines)
diff --git a/server/site_tests/provision_TestbedUpdate/provision_TestbedUpdate.py b/server/site_tests/provision_TestbedUpdate/provision_TestbedUpdate.py
new file mode 100644
index 0000000..a64b719
--- /dev/null
+++ b/server/site_tests/provision_TestbedUpdate/provision_TestbedUpdate.py
@@ -0,0 +1,96 @@
+# 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.server import afe_utils
+from autotest_lib.server import test
+
+
+class provision_TestbedUpdate(test.test):
+    """A test that can provision a machine to the correct Android version."""
+    version = 1
+
+    def _builds_to_set(self, builds):
+        """Helper function to convert a build string into a set of builds.
+
+        @param builds: Testbed build string to convert into a set.
+
+        @returns: A set of the different builds in the build string.
+        """
+        result = set()
+        if not builds:
+            return result
+        builds = builds.split(',')
+        for build in builds:
+            # Remove any build multipliers, i.e. <build>#2
+            build = build.split('#')[0]
+            result.add(build)
+        return result
+
+
+    def initialize(self, host, value, force=False, is_test_na=False,
+                   repair=False):
+        """Initialize.
+
+        @param host: The testbed object to update to |value|.
+                     NOTE: This arg must be called host to align with the other
+                           provision actions.
+        @param value: String of the image we want to install on the testbed.
+        @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.
+        @param repair: not used by initialize.
+        """
+        if is_test_na:
+            raise error.TestNAError('Provisioning not applicable.')
+        # We check value in initialize so that it fails faster.
+        if not (value or repair):
+            raise error.TestFail('No build version specified.')
+
+
+    def run_once(self, host, value=None, force=False, repair=False):
+        """The method called by the control file to start the test.
+
+        @param host: The testbed object to update to |value|.
+                     NOTE: This arg must be called host to align with the other
+                           provision actions.
+        @param value: The testbed 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.
+        @param repair: Not yet supported for testbeds.
+
+        """
+        testbed = host
+        logging.debug('Start provisioning %s to %s', testbed, value)
+
+        if not value and not repair:
+            raise error.TestFail('No build provided and this is not a repair '
+                                 ' job.')
+
+        # If the host is already on the correct build, we have nothing to do.
+        if not force and (self._builds_to_set(afe_utils.get_build(testbed)) ==
+                          self._builds_to_set(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,
+                            'Testbed already running %s' % value)
+            return
+        try:
+            afe_utils.machine_install_and_update_labels(
+                    host, image=value)
+        except error.InstallError as e:
+            logging.exception(e)
+            raise error.TestFail(str(e))
+        logging.debug('Finished provisioning %s to %s', host, value)
\ No newline at end of file
diff --git a/server/site_tests/testbed_DummyTest/control b/server/site_tests/testbed_DummyTest/control
index b8df53f..69142c8 100644
--- a/server/site_tests/testbed_DummyTest/control
+++ b/server/site_tests/testbed_DummyTest/control
@@ -6,6 +6,7 @@
 NAME = 'testbed_DummyTest'
 TIME = 'SHORT'
 TEST_TYPE = 'Server'
+ATTRIBUTES = "suite:dummy_testbed"
 # All android tests do not support server-side packaging.
 REQUIRE_SSP = False