[autotest] Store afe_host in and remove host attributes to/from machine dict.

There are multiple times when the AFE is queried for random bits of info
(host attibutes, labels, etc).  To reduce the number of calls, store the
AFE host object into the machine dict so we can store it into the host class.
I chose the AbstractSSHHost class to grab it out of the arg dict since that
is used both by CrosHost and ADBHost.

Also removed host attributes as a class instance attribute since
afe_host contains all that info already.

Also use the os labels to do host class detection.

BUG=chromium:546741
TEST=Tested a couple scenarios:
- Manually on moblab to make sure cros/adb/testbed devices all pass
  dummy tests.
- Trybot job on gnawty-paladin.
- Moblab dummy suite run to ensure no breaking backward compatibility changes
  (moblab on R52 and testing with a tryjob gnawty build).
- test_{that,droid}.py on a cros/adbhost device.

Change-Id: I101c478c0f27f9fd444ccd4699a5c94e76d3b17e
Reviewed-on: https://chromium-review.googlesource.com/349990
Commit-Ready: Kevin Cheng <kevcheng@chromium.org>
Tested-by: Kevin Cheng <kevcheng@chromium.org>
Reviewed-by: Kevin Cheng <kevcheng@chromium.org>
diff --git a/server/hosts/abstract_ssh.py b/server/hosts/abstract_ssh.py
index 4ee54c3..0dd294e 100644
--- a/server/hosts/abstract_ssh.py
+++ b/server/hosts/abstract_ssh.py
@@ -26,10 +26,19 @@
     VERSION_PREFIX = ''
 
     def _initialize(self, hostname, user="root", port=22, password="",
-                    is_client_install_supported=True, host_attributes={},
+                    is_client_install_supported=True, afe_host=None,
                     *args, **dargs):
         super(AbstractSSHHost, self)._initialize(hostname=hostname,
                                                  *args, **dargs)
+        """
+        @param hostname: The hostname of the host.
+        @param user: The username to use when ssh'ing into the host.
+        @param password: The password to use when ssh'ing into the host.
+        @param port: The port to use for ssh.
+        @param is_client_install_supported: Boolean to indicate if we can
+                install autotest on the host.
+        @param afe_host: The host object attained from the AFE (get_hosts).
+        """
         # IP address is retrieved only on demand. Otherwise the host
         # initialization will fail for host is not online.
         self._ip = None
@@ -53,8 +62,7 @@
         # Create a Lock to protect against race conditions.
         self._lock = Lock()
 
-        self.host_attributes = host_attributes
-
+        self._afe_host = afe_host or utils.EmptyAFEHost()
 
     @property
     def ip(self):
diff --git a/server/hosts/adb_host.py b/server/hosts/adb_host.py
index 5e157df..857cf8f 100644
--- a/server/hosts/adb_host.py
+++ b/server/hosts/adb_host.py
@@ -218,11 +218,10 @@
 
         self.tmp_dirs = []
         self.labels = base_label.LabelRetriever(adb_label.ADB_LABELS)
-        # TODO (sbasi/kevcheng): Once the teststation host is committed,
-        # refactor the serial retrieval.
-        adb_serial = adb_serial or self.host_attributes.get('serials', None)
-        fastboot_serial = fastboot_serial or self.host_attributes.get(
-                'fastboot_serial', None)
+        adb_serial = adb_serial or self._afe_host.attributes.get('serials')
+        fastboot_serial = (fastboot_serial or
+                self._afe_host.attributes.get('fastboot_serial'))
+
         self.adb_serial = adb_serial
         if adb_serial:
             adb_prefix = any(adb_serial.startswith(p)
diff --git a/server/hosts/base_label.py b/server/hosts/base_label.py
index bdf269b..26bc2c4 100644
--- a/server/hosts/base_label.py
+++ b/server/hosts/base_label.py
@@ -232,8 +232,7 @@
             self._populate_known_labels(self._labels)
 
         afe = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10)
-        afe_host = afe.get_hosts(hostname=host.hostname)[0]
-        old_labels = set(afe_host.labels)
+        old_labels = set(host._afe_host.labels)
         logging.info('existing labels: %s', old_labels)
         known_labels = set([l for l in old_labels
                             if self._is_known_label(l)])
diff --git a/server/hosts/base_label_unittest.py b/server/hosts/base_label_unittest.py
index d6fdf37..cdcab9f 100644
--- a/server/hosts/base_label_unittest.py
+++ b/server/hosts/base_label_unittest.py
@@ -8,6 +8,7 @@
 import common
 
 from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
+from autotest_lib.server import utils
 from autotest_lib.server.hosts import base_label
 
  # pylint: disable=missing-docstring
@@ -64,11 +65,18 @@
         return labels
 
 
+class MockAFEHost(utils.EmptyAFEHost):
+
+    def __init__(self, labels=[]):
+        self.labels = labels
+
+
 class MockHost(object):
 
-    def __init__(self, exists=True):
+    def __init__(self, exists=True, afe_host=None):
         self.hostname = 'hostname'
         self.exists = exists
+        self._afe_host = afe_host
 
 
 class BaseLabelUnittests(unittest.TestCase):
@@ -178,12 +186,9 @@
         """Check that we add/remove the expected labels in update_labels()."""
         label_to_add = 'label_to_add'
         label_to_remove = 'prefix:label_to_remove'
-        mock_afe_host = mock.MagicMock()
-        mock_afe_host.labels = [label_to_remove, TestBaseLabel._NAME]
         mock_afe = mock.MagicMock()
-        mock_afe.get_hosts.return_value = [mock_afe_host]
-        mock_retry_afe.return_value = mock_afe
-        mockhost = MockHost()
+        mockhost = MockHost(afe_host=MockAFEHost(
+                labels=[label_to_remove, TestBaseLabel._NAME]))
         expected_remove_labels = [label_to_remove]
         expected_add_labels = ['%s:%s' % (TestStringPrefixLabel._NAME,
                                           label_to_add)]
diff --git a/server/hosts/factory.py b/server/hosts/factory.py
index be85453..d165a97 100644
--- a/server/hosts/factory.py
+++ b/server/hosts/factory.py
@@ -7,7 +7,7 @@
 from autotest_lib.client.bin import utils
 from autotest_lib.client.common_lib import error, global_config
 from autotest_lib.server import utils as server_utils
-from autotest_lib.server.hosts import cros_host, ssh_host
+from autotest_lib.server.hosts import common_label, cros_host, ssh_host
 from autotest_lib.server.hosts import moblab_host, sonic_host
 from autotest_lib.server.hosts import adb_host, testbed
 
@@ -32,7 +32,8 @@
 host_types = [cros_host.CrosHost, moblab_host.MoblabHost, sonic_host.SonicHost,
               adb_host.ADBHost,]
 OS_HOST_DICT = {'cros' : cros_host.CrosHost,
-                'android': adb_host.ADBHost}
+                'android': adb_host.ADBHost,
+                'brillo': adb_host.ADBHost}
 
 
 def _get_host_arguments():
@@ -121,7 +122,7 @@
     @param machine: A dict representing the device under test or a String
                     representing the DUT 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
+                    Optional 'afe_host' key will pipe in afe_host
                     from the autoserv runtime or the AFE.
     @param host_class: Host class to use, if None, will attempt to detect
                        the correct class.
@@ -133,9 +134,8 @@
     @returns: A host object which is an instance of the newly created
               host class.
     """
-    hostname, host_attributes = server_utils.get_host_info_from_machine(
+    hostname, afe_host = server_utils.get_host_info_from_machine(
             machine)
-    args['host_attributes'] = host_attributes
     ssh_user, ssh_pass, ssh_port, ssh_verbosity_flag, ssh_options = \
             _get_host_arguments()
 
@@ -144,11 +144,22 @@
     args['ssh_verbosity_flag'] = ssh_verbosity_flag
     args['ssh_options'] = ssh_options
     args['port'] = ssh_port
+    args['afe_host'] = afe_host
+
+    host_os = None
+    # Let's grab the os from the labels if we can for host class detection.
+    for label in afe_host.labels:
+        if label.startswith(common_label.OSLabel._NAME):
+            host_os = label.split(':')[-1]
+            break
 
     if not connectivity_class:
         connectivity_class = _choose_connectivity_class(hostname, ssh_port)
-    host_attributes = args.get('host_attributes', {})
-    host_class = host_class or OS_HOST_DICT.get(host_attributes.get('os_type'))
+    # TODO(kevcheng): get rid of the host detection using host attributes.
+    host_class = (host_class
+                  or OS_HOST_DICT.get(afe_host.attributes.get('os_type'))
+                  or OS_HOST_DICT.get(host_os))
+
     if host_class:
         classes = [host_class, connectivity_class]
     else:
@@ -174,15 +185,15 @@
                     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.
+                    Optional 'afe_host' key will pipe in afe_host from
+                    the afe_host object from the autoserv runtime or the AFE.
     @param kwargs: Keyword args to pass to the testbed initialization.
 
     @returns: The testbed object with all associated host objects instantiated.
     """
-    hostname, host_attributes = server_utils.get_host_info_from_machine(
+    hostname, afe_host = server_utils.get_host_info_from_machine(
             machine)
-    kwargs['host_attributes'] = host_attributes
+    kwargs['afe_host'] = afe_host
     return testbed.TestBed(hostname, **kwargs)
 
 
@@ -193,7 +204,7 @@
                     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
+                    Optional 'afe_host' key will pipe in afe_host
                     from the autoserv runtime or the AFE.
     @param kwargs: Keyword args to pass to the testbed initialization.
 
diff --git a/server/hosts/servo_host.py b/server/hosts/servo_host.py
index fca3243..0aacce9 100644
--- a/server/hosts/servo_host.py
+++ b/server/hosts/servo_host.py
@@ -732,7 +732,7 @@
 
     This checks for the presence of servo host and port attached to the
     given `dut_host`.  This data should be stored in the
-    `host_attributes` field in the provided `dut_host` parameter.
+    `_afe_host.attributes` field in the provided `dut_host` parameter.
 
     @param dut_host   Instance of `Host` on which to find the servo
                       attributes.
@@ -749,7 +749,7 @@
         is_ssp_moblab = is_moblab
     else:
         is_moblab = utils.is_moblab()
-    attrs = dut_host.host_attributes
+    attrs = dut_host._afe_host.attributes
     if attrs and SERVO_HOST_ATTR in attrs:
         servo_host = attrs[SERVO_HOST_ATTR]
         if (is_ssp_moblab and servo_host in ['localhost', '127.0.0.1']):
diff --git a/server/hosts/testbed.py b/server/hosts/testbed.py
index 68de62b..ae67baa 100644
--- a/server/hosts/testbed.py
+++ b/server/hosts/testbed.py
@@ -14,8 +14,8 @@
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.common_lib import logging_config
 from autotest_lib.server.cros.dynamic_suite import constants
-from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
 from autotest_lib.server import autoserv_parser
+from autotest_lib.server import utils
 from autotest_lib.server.hosts import adb_host
 from autotest_lib.server.hosts import base_label
 from autotest_lib.server.hosts import testbed_label
@@ -38,25 +38,25 @@
     _parser = autoserv_parser.autoserv_parser
     VERSION_PREFIX = 'testbed-version'
 
-    def __init__(self, hostname='localhost', host_attributes={},
-                 adb_serials=None, **dargs):
+    def __init__(self, hostname='localhost', afe_host=None, adb_serials=None,
+                 **dargs):
         """Initialize a TestBed.
 
         This will create the Test Station Host and connected hosts (ADBHost for
         now) and allow the user to retrieve them.
 
         @param hostname: Hostname of the test station connected to the duts.
-        @param host_attributes: Attributes of the host, passed in from
-                factory.create_testbed.
         @param adb_serials: List of adb device serials.
+        @param afe_host: The host object attained from the AFE (get_hosts).
         """
         logging.info('Initializing TestBed centered on host: %s', hostname)
         self.hostname = hostname
+        self._afe_host = afe_host or utils.EmptyAFEHost()
         self.labels = base_label.LabelRetriever(testbed_label.TESTBED_LABELS)
         self.teststation = teststation_host.create_teststationhost(
                 hostname=hostname)
         self.is_client_install_supported = False
-        serials_from_attributes = host_attributes.get('serials')
+        serials_from_attributes = self._afe_host.attributes.get('serials')
         if serials_from_attributes:
             serials_from_attributes = serials_from_attributes.split(',')
 
@@ -69,7 +69,6 @@
                 hostname=hostname, teststation=self.teststation,
                 adb_serial=adb_serial)
 
-        self.host_attributes = host_attributes
 
 
     def query_adb_device_serials(self):
@@ -77,23 +76,9 @@
 
         @returns a list of adb devices.
         """
-        serials = []
-        # Let's see if we can get the serials via host attributes.
-        afe = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10)
-        serials_attr = afe.get_host_attribute('serials', hostname=self.hostname)
-        for serial_attr in serials_attr:
-            serials.extend(serial_attr.value.split(','))
-
-        # Looks like we got nothing from afe, let's probe the test station.
-        if not serials:
-            # TODO(kevcheng): Refactor teststation to be a class and make the
-            # ADBHost adb_devices a static method I can use here.  For now this
-            # is pretty much a c/p of the _adb_devices() method from ADBHost.
-            serials = adb_host.ADBHost.parse_device_serials(
+        return adb_host.ADBHost.parse_device_serials(
                 self.teststation.run('adb devices').stdout)
 
-        return serials
-
 
     def get_all_hosts(self):
         """Return a list of all the hosts in this testbed.