[autotest] Update Repair and Verify to handle a testbed.

The repair and verify flow assumes the machine passed in is a host.
Both control segments now check to see if it's dealing with a testbed
and to run through a list of hosts and run the proper host methods to
ensure a successful verify/repair.

BUG=chromium:548739
TEST=Locally ran reverify/repair on testbed with shamu/angler/hammerhead
successfully.

Change-Id: Ib0ade22dd9e3d7496c083fe48b642ed4069eadd6
Reviewed-on: https://chromium-review.googlesource.com/318385
Commit-Ready: Kevin Cheng <kevcheng@chromium.org>
Tested-by: Kevin Cheng <kevcheng@chromium.org>
Reviewed-by: Richard Barnette <jrbarnette@chromium.org>
Reviewed-by: Simran Basi <sbasi@chromium.org>
diff --git a/server/control_segments/repair b/server/control_segments/repair
index 0ba5966..da379f6 100644
--- a/server/control_segments/repair
+++ b/server/control_segments/repair
@@ -15,17 +15,23 @@
     try:
         hostname = utils.get_hostname_from_machine(machine)
         job.record('START', None, 'repair')
-        host = hosts.create_host(machine, initialize=False, auto_monitor=False,
-                                 try_lab_servo=True)
-        # Collect logs before the repair, as it might destroy all useful logs.
-        local_log_dir = os.path.join(job.resultdir, hostname, 'before_repair')
-        host.collect_logs('/var/log', local_log_dir, ignore_errors=True)
-        # Collect crash info.
-        crashcollect.get_crashinfo(host, None)
+        target = hosts.create_target_machine(machine, initialize=False,
+                                             auto_monitor=False,
+                                             try_lab_servo=True)
+        # We don't need to collect logs or crash info if we're a testbed since
+        # they're not applicable (yet).
+        if not utils.machine_is_testbed(machine):
+            # Collect logs before the repair, as it might destroy all
+            # useful logs.
+            local_log_dir = os.path.join(job.resultdir, hostname,
+                                         'before_repair')
+            target.collect_logs('/var/log', local_log_dir, ignore_errors=True)
+            # Collect crash info.
+            crashcollect.get_crashinfo(target, None)
 
-        host.repair()
+        target.repair()
         logging.debug('Repair with labels list %s', labels_list)
-        provision.run_special_task_actions(job, host, labels_list,
+        provision.run_special_task_actions(job, target, labels_list,
                                            provision.Repair)
     except Exception as e:
         logging.exception(e)
diff --git a/server/control_segments/verify b/server/control_segments/verify
index 9e9d6f3..a86fc83 100644
--- a/server/control_segments/verify
+++ b/server/control_segments/verify
@@ -1,4 +1,5 @@
 from autotest_lib.client.common_lib.cros.graphite import autotest_stats
+from autotest_lib.server import utils
 from autotest_lib.server.cros import provision
 
 
@@ -10,17 +11,17 @@
 def verify(machine):
     print 'Initializing host %s' % machine
     timer = None
+    # We set try_lab_servo to True to trigger servo verify and
+    # servo update if needed.
+    target = hosts.create_target_machine(machine, initialize=False,
+                                         auto_monitor=False, try_lab_servo=True)
     try:
         job.record('START', None, 'verify')
-        # We set try_lab_servo to True to trigger servo verify and
-        # servo update if needed.
-        host = hosts.create_host(machine, initialize=False, auto_monitor=False,
-                                 try_lab_servo=True)
         timer = autotest_stats.Timer('verify_time')
         timer.start()
 
-        host.verify()
-        provision.run_special_task_actions(job, host, labels_list,
+        target.verify()
+        provision.run_special_task_actions(job, target, labels_list,
                                            provision.Verify)
     except Exception as e:
         logging.exception(e)
diff --git a/server/hosts/__init__.py b/server/hosts/__init__.py
index ba70e29..e16dc5d 100644
--- a/server/hosts/__init__.py
+++ b/server/hosts/__init__.py
@@ -27,3 +27,4 @@
 # factory function
 from factory import create_host
 from factory import create_testbed
+from factory import create_target_machine
diff --git a/server/hosts/factory.py b/server/hosts/factory.py
index dd20a6e..458f77e 100644
--- a/server/hosts/factory.py
+++ b/server/hosts/factory.py
@@ -180,4 +180,24 @@
     hostname, host_attributes = server_utils.get_host_info_from_machine(
             machine)
     kwargs['host_attributes'] = host_attributes
-    return testbed.TestBed(hostname, kwargs)
+    return testbed.TestBed(hostname, **kwargs)
+
+
+def create_target_machine(machine, **kwargs):
+    """Create the target machine which could be a testbed or a *Host.
+
+    @param machine: A dict representing the test bed under test or a String
+                    representing the testbed hostname (for legacy caller
+                    support).
+                    If it is a machine dict, the 'hostname' key is required.
+                    Optional 'host_attributes' key will pipe in host_attributes
+                    from the autoserv runtime or the AFE.
+    @param kwargs: Keyword args to pass to the testbed initialization.
+
+    @returns: The target machine to be used for verify/repair.
+    """
+    # TODO(kevcheng): We'll want to have a smarter way of figuring out which
+    # host to create (checking host labels).
+    if server_utils.machine_is_testbed(machine):
+        return create_testbed(machine, **kwargs)
+    return create_host(machine, **kwargs)
diff --git a/server/hosts/testbed.py b/server/hosts/testbed.py
index e1fbb4e..42604ca 100644
--- a/server/hosts/testbed.py
+++ b/server/hosts/testbed.py
@@ -19,7 +19,7 @@
 
 
     def __init__(self, hostname='localhost', host_attributes={},
-                 adb_serials=None):
+                 adb_serials=None, **dargs):
         """Initialize a TestBed.
 
         This will create the Test Station Host and connected hosts (ADBHost for
@@ -34,7 +34,7 @@
                 hostname=hostname)
         serials_from_attributes = host_attributes.get('serials')
         if serials_from_attributes:
-            serials_from_attributes.split(',')
+            serials_from_attributes = serials_from_attributes.split(',')
 
         self.adb_device_serials = (adb_serials or
                                    serials_from_attributes or
@@ -76,7 +76,8 @@
                  devices.
         """
         device_list = [self.teststation]
-        return device_list.extend(self.adb_devices.values())
+        device_list.extend(self.adb_devices.values())
+        return device_list
 
 
     def get_test_station(self):
@@ -130,3 +131,15 @@
         @return: A string representing the testbed platform.
         """
         return 'testbed'
+
+
+    def repair(self):
+        """Run through repair on all the devices."""
+        for adb_device in self.get_adb_devices().values():
+            adb_device.repair()
+
+
+    def verify(self):
+        """Run through verify on all the devices."""
+        for device in self.get_all_hosts():
+            device.verify()
diff --git a/server/site_utils.py b/server/site_utils.py
index 52ab791..623c92e 100644
--- a/server/site_utils.py
+++ b/server/site_utils.py
@@ -718,3 +718,19 @@
     if not creds_dir or not os.path.exists(creds_dir):
         creds_dir = common.autotest_dir
     return os.path.join(creds_dir, creds_file)
+
+
+def machine_is_testbed(machine):
+    """Checks if the machine is a testbed.
+
+    The signal we use to determine if the machine is a testbed
+    is if the host attributes contain more than 1 serial.
+
+    @param machine: is a list of dicts
+
+    @return: True if the machine is a testbed, False otherwise.
+    """
+    _, attributes = get_host_info_from_machine(machine)
+    if len(attributes.get('serials', '').split(',')) > 1:
+        return True
+    return False