[autotest] autoserv add --lab & --host_attributes arguments

Added two new flags to autoserv.

--lab indicates that autoserv is running in the lab and has
the full Autotest infrastructure at its disposal.

--host_attributes allows host attribute information that is
usually in the database to be retrievable from the command
line arguments.

If --lab is pulled in, autoserv will request the host attributes
from the database at test runtime.

From here this change, then updates the concept of the "machines"
list that test control files receive to now be a list of dicts
that contain the machine hostname and host attributes. This will
enable identifing information the hosts library needs to create
host objects to be available whether or not there is a database
present.

BUG=chromium:564343
TEST=local autoserv runs. Also verified scheduler changes work
via MobLab. waiting on trybot results.
DEPLOY=scheduler

Change-Id: I6021de11317e29e2e6c084d863405910c7d1a71d
Reviewed-on: https://chromium-review.googlesource.com/315230
Commit-Ready: Simran Basi <sbasi@chromium.org>
Tested-by: Simran Basi <sbasi@chromium.org>
Reviewed-by: Simran Basi <sbasi@chromium.org>
diff --git a/client/bin/job.py b/client/bin/job.py
index 11a4175..20bb276 100644
--- a/client/bin/job.py
+++ b/client/bin/job.py
@@ -5,16 +5,15 @@
 Copyright Andy Whitcroft, Martin J. Bligh 2006
 """
 
-import copy, os, platform, re, shutil, sys, time, traceback, types, glob
-import logging, getpass, errno, weakref
-import cPickle as pickle
+import copy, os, re, shutil, sys, time, traceback, types, glob
+import logging, getpass, weakref
 from autotest_lib.client.bin import client_logging_config
 from autotest_lib.client.bin import utils, parallel, kernel, xen
 from autotest_lib.client.bin import profilers, boottool, harness
 from autotest_lib.client.bin import config, sysinfo, test, local_host
 from autotest_lib.client.bin import partition as partition_lib
 from autotest_lib.client.common_lib import base_job
-from autotest_lib.client.common_lib import error, barrier, log, logging_manager
+from autotest_lib.client.common_lib import error, barrier, logging_manager
 from autotest_lib.client.common_lib import base_packages, packages
 from autotest_lib.client.common_lib import global_config
 from autotest_lib.client.tools import html_report
@@ -245,6 +244,10 @@
         self._init_bootloader()
 
         self.machines = [options.hostname]
+        self.machine_dict_list = [{'hostname' : options.hostname}]
+        # Client side tests should always run the same whether or not they are
+        # running in the lab.
+        self.in_lab = False
         self.hosts = set([local_host.LocalHost(hostname=options.hostname,
                                                bootloader=self.bootloader)])
 
diff --git a/client/bin/setup_job.py b/client/bin/setup_job.py
index e87d8f8..b8de3a4 100644
--- a/client/bin/setup_job.py
+++ b/client/bin/setup_job.py
@@ -5,7 +5,6 @@
 import logging, os, pickle, re, sys
 import common
 
-from autotest_lib.client.bin import client_logging_config
 from autotest_lib.client.bin import job as client_job
 from autotest_lib.client.common_lib import base_job
 from autotest_lib.client.common_lib import error
@@ -32,6 +31,10 @@
         base_job.base_job.__init__(self, options=options)
         self._cleanup_debugdir_files()
         self._cleanup_results_dir()
+        self.machine_dict_list = [{'hostname' : options.hostname}]
+        # Client side tests should always run the same whether or not they are
+        # running in the lab.
+        self.in_lab = False
         self.pkgmgr = packages.PackageManager(
             self.autodir, run_function_dargs={'timeout':3600})
 
diff --git a/client/common_lib/base_job_unittest.py b/client/common_lib/base_job_unittest.py
index 9f84e17..acd6087 100755
--- a/client/common_lib/base_job_unittest.py
+++ b/client/common_lib/base_job_unittest.py
@@ -83,7 +83,7 @@
             'num_tests_run', 'pkgmgr', 'profilers', 'resultdir',
             'run_test_cleanup', 'sysinfo', 'tag', 'user', 'use_sequence_number',
             'warning_loggers', 'warning_manager', 'label', 'test_retry',
-            'parent_job_id'
+            'parent_job_id', 'in_lab', 'machine_dict_list'
             ])
 
         OPTIONAL_ATTRIBUTES = set([
diff --git a/scheduler/monitor_db.py b/scheduler/monitor_db.py
index 77e1b55..f7a9dc4 100755
--- a/scheduler/monitor_db.py
+++ b/scheduler/monitor_db.py
@@ -263,10 +263,11 @@
     @param queue_entry - A HostQueueEntry object - If supplied and no Job
             object was supplied, this will be used to lookup the Job object.
     """
-    return autoserv_utils.autoserv_run_job_command(_autoserv_directory,
+    command = autoserv_utils.autoserv_run_job_command(_autoserv_directory,
             machines, results_directory=drone_manager.WORKING_DIRECTORY,
             extra_args=extra_args, job=job, queue_entry=queue_entry,
-            verbose=verbose)
+            verbose=verbose, in_lab=True)
+    return command
 
 
 class BaseDispatcher(object):
diff --git a/scheduler/monitor_db_unittest.py b/scheduler/monitor_db_unittest.py
index 56a04e5..66cdc7c 100755
--- a/scheduler/monitor_db_unittest.py
+++ b/scheduler/monitor_db_unittest.py
@@ -1071,6 +1071,7 @@
         extra_args = ['-Z', 'hello']
         expected_command_line_base = set((monitor_db._autoserv_path, '-p',
                                           '-m', machines, '-r',
+                                          '--lab', 'True',
                                           drone_manager.WORKING_DIRECTORY))
 
         expected_command_line = expected_command_line_base.union(
diff --git a/server/autoserv b/server/autoserv
index 8e36aec..889ff07 100755
--- a/server/autoserv
+++ b/server/autoserv
@@ -442,6 +442,8 @@
     ssh_verbosity = int(parser.options.ssh_verbosity)
     ssh_options = parser.options.ssh_options
     no_use_packaging = parser.options.no_use_packaging
+    host_attributes = parser.options.host_attributes
+    in_lab = bool(parser.options.lab)
 
     # can't be both a client and a server side test
     if client and server:
@@ -478,6 +480,9 @@
         kwargs['parent_job_id'] = int(parser.options.parent_job_id)
     if control_filename:
         kwargs['control_filename'] = control_filename
+    if host_attributes:
+        kwargs['host_attributes'] = host_attributes
+    kwargs['in_lab'] = in_lab
     job = server_job.server_job(control, parser.args[1:], results, label,
                                 user, machines, client, parse_job,
                                 ssh_user, ssh_port, ssh_pass,
diff --git a/server/autoserv_parser.py b/server/autoserv_parser.py
index 0f5c3ea..82f6689 100644
--- a/server/autoserv_parser.py
+++ b/server/autoserv_parser.py
@@ -1,4 +1,5 @@
 import argparse
+import ast
 import logging
 import os
 import shlex
@@ -201,6 +202,16 @@
                                      'R27-3837.0.0 or a build name: '
                                      'x86-alex-release/R27-3837.0.0 to '
                                      'utilize lab devservers automatically.'))
+        self.parser.add_argument('--host_attributes', action='store',
+                                 dest='host_attributes', default='{}',
+                                 help=('Host attribute to be applied to all '
+                                       'machines/hosts for this autoserv run. '
+                                       'Must be a string-encoded dict. '
+                                       'Example: {"key1":"value1", "key2":'
+                                       '"value2"}'))
+        self.parser.add_argument('--lab', action='store', type=str,
+                                 dest='lab', default='',
+                                 help=argparse.SUPPRESS)
         #
         # Warning! Please read before adding any new arguments!
         #
@@ -239,6 +250,8 @@
         if self.options.image:
             self.options.install_before = True
             self.options.image =  self.options.image.strip()
+        self.options.host_attributes = ast.literal_eval(
+                self.options.host_attributes)
 
 
 # create the one and only one instance of autoserv_parser
diff --git a/server/autoserv_utils.py b/server/autoserv_utils.py
index 3275fb4..3a11cca 100644
--- a/server/autoserv_utils.py
+++ b/server/autoserv_utils.py
@@ -28,7 +28,8 @@
                              ssh_verbosity=0,
                              no_console_prefix=False,
                              ssh_options=None,
-                             use_packaging=True):
+                             use_packaging=True,
+                             in_lab=False):
     """
     Construct an autoserv command from a job or host queue entry.
 
@@ -56,6 +57,10 @@
     @param ssh_options: A string giving extra arguments to be tacked on to
                         ssh commands.
     @param use_packaging Enable install modes that use the packaging system.
+    @param in_lab: If true, informs autoserv it is running within a lab
+                   environment. This information is useful as autoserv knows
+                   the database is available and can make database calls such
+                   as looking up host attributes at runtime.
 
     @returns The autoserv command line as a list of executable + parameters.
 
@@ -114,6 +119,9 @@
     if not use_packaging:
         command.append('--no_use_packaging')
 
+    if in_lab:
+        command.extend(['--lab', 'True'])
+
     return command + extra_args
 
 
diff --git a/server/control_segments/cleanup b/server/control_segments/cleanup
index d8b585c..ecdb20d 100644
--- a/server/control_segments/cleanup
+++ b/server/control_segments/cleanup
@@ -14,6 +14,7 @@
 def cleanup(machine):
     timer = None
     try:
+        hostname = utils.get_hostname_from_machine(machine)
         job.record('START', None, 'cleanup')
         host = hosts.create_host(machine, initialize=False, auto_monitor=False,
                                  try_lab_servo=True)
@@ -30,7 +31,7 @@
         except error.AutoservSshPingHostError:
             # Try to restart dut with servo.
             host._servo_repair_power()
-        local_log_dir = os.path.join(job.resultdir, machine)
+        local_log_dir = os.path.join(job.resultdir, hostname)
         host.collect_logs('/var/log', local_log_dir, ignore_errors=True)
 
         host.cleanup()
diff --git a/server/control_segments/repair b/server/control_segments/repair
index 3d740a2..e2c2161 100644
--- a/server/control_segments/repair
+++ b/server/control_segments/repair
@@ -12,11 +12,12 @@
 
 def repair(machine):
     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, machine, 'before_repair')
+        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)
diff --git a/server/control_segments/reset b/server/control_segments/reset
index ebe01d1..c03369a 100644
--- a/server/control_segments/reset
+++ b/server/control_segments/reset
@@ -10,7 +10,7 @@
 
 
 def reset(machine):
-    print 'Starting to reset host ' + machine
+    print 'Starting to reset host %s' % machine
     timer = None
     try:
         job.record('START', None, 'reset')
diff --git a/server/control_segments/verify b/server/control_segments/verify
index 268a5e7..9e9d6f3 100644
--- a/server/control_segments/verify
+++ b/server/control_segments/verify
@@ -8,7 +8,7 @@
 
 
 def verify(machine):
-    print 'Initializing host ' + machine
+    print 'Initializing host %s' % machine
     timer = None
     try:
         job.record('START', None, 'verify')
diff --git a/server/hosts/abstract_ssh.py b/server/hosts/abstract_ssh.py
index 6a7367f..d8f51b3 100644
--- a/server/hosts/abstract_ssh.py
+++ b/server/hosts/abstract_ssh.py
@@ -24,7 +24,8 @@
     """
 
     def _initialize(self, hostname, user="root", port=22, password="",
-                    is_client_install_supported=True, *args, **dargs):
+                    is_client_install_supported=True, host_attributes={},
+                    *args, **dargs):
         super(AbstractSSHHost, self)._initialize(hostname=hostname,
                                                  *args, **dargs)
         # IP address is retrieved only on demand. Otherwise the host
@@ -50,6 +51,8 @@
         # Create a Lock to protect against race conditions.
         self._lock = Lock()
 
+        self.host_attributes = host_attributes
+
 
     @property
     def ip(self):
diff --git a/server/hosts/adb_host.py b/server/hosts/adb_host.py
index 1489e76..a946b3d 100644
--- a/server/hosts/adb_host.py
+++ b/server/hosts/adb_host.py
@@ -153,6 +153,10 @@
                                 run over TCP/IP.
         @param teststation: The teststation object ADBHost should use.
         """
+        # Sets up the is_client_install_supported field.
+        super(ADBHost, self)._initialize(hostname=hostname,
+                                         is_client_install_supported=False,
+                                         *args, **dargs)
         if device_hostname and (adb_serial or fastboot_serial):
             raise error.AutoservError(
                     'TCP/IP and USB modes are mutually exclusive')
@@ -163,6 +167,9 @@
         self._num_of_boards = {}
         self._device_hostname = device_hostname
         self._use_tcpip = False
+        # TODO (sbasi/kevcheng): Once the teststation host is committed,
+        # refactor the serial retrieval.
+        adb_serial = adb_serial or self.host_attributes.get('serials', None)
         self._adb_serial = adb_serial
         self._fastboot_serial = fastboot_serial or adb_serial
         self.teststation = teststation
@@ -180,11 +187,6 @@
             msg += ', fastboot serial: %s' % self._fastboot_serial
         logging.debug(msg)
 
-        # Sets up the is_client_install_supported field.
-        super(ADBHost, self)._initialize(hostname=hostname,
-                                         is_client_install_supported=False,
-                                         *args, **dargs)
-
         # Make sure teststation credentials work.
         try:
             self.teststation.run('true')
diff --git a/server/hosts/factory.py b/server/hosts/factory.py
index 1029506..a841814 100644
--- a/server/hosts/factory.py
+++ b/server/hosts/factory.py
@@ -81,13 +81,17 @@
     return cros_host.CrosHost
 
 
-def create_host(hostname, host_class=None, **args):
+def create_host(machine, host_class=None, **args):
     """Create a host object.
 
     This method mixes host classes that are needed into a new subclass
     and creates a instance of the new class.
 
-    @param hostname: A string representing the host name of the device.
+    @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
+                    from the autoserv runtime or the AFE.
     @param host_class: Host class to use, if None, will attempt to detect
                        the correct class.
     @param args: Args that will be passed to the constructor of
@@ -96,6 +100,9 @@
     @returns: A host object which is an instance of the newly created
               host class.
     """
+    hostname, host_attributes = 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()
 
diff --git a/server/hosts/sonic_host.py b/server/hosts/sonic_host.py
index b593906..9b247b3 100644
--- a/server/hosts/sonic_host.py
+++ b/server/hosts/sonic_host.py
@@ -58,7 +58,8 @@
         """
         try:
             result = host.run('getprop ro.product.device', timeout=timeout)
-        except (error.AutoservRunError, error.AutoservSSHTimeout):
+        except (error.AutoservRunError, error.AutoservSSHTimeout,
+                error.AutotestHostRunError):
             return False
         return 'anchovy' in result.stdout
 
diff --git a/server/server_job.py b/server/server_job.py
index 17c6790..92bf33e 100644
--- a/server/server_job.py
+++ b/server/server_job.py
@@ -18,6 +18,8 @@
 from autotest_lib.client.common_lib import error, utils, packages
 from autotest_lib.client.common_lib import logging_manager
 from autotest_lib.server import test, subcommand, profilers
+from autotest_lib.server import utils as server_utils
+from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
 from autotest_lib.server.hosts import abstract_ssh, factory as host_factory
 from autotest_lib.tko import db as tko_db, status_lib, utils as tko_utils
 
@@ -153,7 +155,7 @@
                  test_retry=0, group_name='',
                  tag='', disable_sysinfo=False,
                  control_filename=SERVER_CONTROL_FILENAME,
-                 parent_job_id=None):
+                 parent_job_id=None, host_attributes=None, in_lab=False):
         """
         Create a server side job object.
 
@@ -183,6 +185,11 @@
                 should be written in the results directory.
         @param parent_job_id: Job ID of the parent job. Default to None if the
                 job does not have a parent job.
+        @param host_attributes: Dict of host attributes passed into autoserv
+                                via the command line. If specified here, these
+                                attributes will apply to all machines.
+        @param in_lab: Boolean that indicates if this is running in the lab
+                       environment.
         """
         super(base_server_job, self).__init__(resultdir=resultdir,
                                               test_retry=test_retry)
@@ -268,6 +275,17 @@
         self.failed_with_device_error = False
 
         self.parent_job_id = parent_job_id
+        self.in_lab = in_lab
+        afe = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10)
+        self.machine_dict_list = []
+        for machine in self.machines:
+            host_attributes = host_attributes or {}
+            if self.in_lab:
+                host = afe.get_hosts(hostname=machine)[0]
+                host_attributes.update(host.attributes)
+            self.machine_dict_list.append(
+                    {'hostname' : machine,
+                     'host_attributes' : host_attributes})
 
 
     @classmethod
@@ -374,7 +392,7 @@
         machines, job, ssh_user, ssh_port, ssh_pass, ssh_verbosity_flag,
         and ssh_options.
         """
-        namespace = {'machines' : self.machines,
+        namespace = {'machines' : self.machine_dict_list,
                      'job' : self,
                      'ssh_user' : self._ssh_user,
                      'ssh_port' : self._ssh_port,
@@ -493,21 +511,23 @@
         is_forking = not (len(machines) == 1 and self.machines == machines)
         if self._parse_job and is_forking and log:
             def wrapper(machine):
-                self._parse_job += "/" + machine
+                hostname = server_utils.get_hostname_from_machine(machine)
+                self._parse_job += "/" + hostname
                 self._using_parser = INCREMENTAL_TKO_PARSING
                 self.machines = [machine]
-                self.push_execution_context(machine)
+                self.push_execution_context(hostname)
                 os.chdir(self.resultdir)
-                utils.write_keyval(self.resultdir, {"hostname": machine})
+                utils.write_keyval(self.resultdir, {"hostname": hostname})
                 self.init_parser()
                 result = function(machine)
                 self.cleanup_parser()
                 return result
         elif len(machines) > 1 and log:
             def wrapper(machine):
-                self.push_execution_context(machine)
+                hostname = server_utils.get_hostname_from_machine(machine)
+                self.push_execution_context(hostname)
                 os.chdir(self.resultdir)
-                machine_data = {'hostname' : machine,
+                machine_data = {'hostname' : hostname,
                                 'status_version' : str(self._STATUS_VERSION)}
                 utils.write_keyval(self.resultdir, machine_data)
                 result = function(machine)
diff --git a/server/server_job_unittest.py b/server/server_job_unittest.py
index db0cd4d..dd7d51f 100755
--- a/server/server_job_unittest.py
+++ b/server/server_job_unittest.py
@@ -1,6 +1,5 @@
 #!/usr/bin/python
 
-import tempfile
 import common
 
 from autotest_lib.server import server_job
@@ -33,7 +32,7 @@
     OPTIONAL_ATTRIBUTES = (
         base_job_unittest.test_init.generic_tests.OPTIONAL_ATTRIBUTES
         - set(['serverdir', 'num_tests_run', 'num_tests_failed',
-               'warning_manager', 'warning_loggers']))
+               'warning_manager', 'warning_loggers', 'in_lab']))
 
     def setUp(self):
         self.god = mock.mock_god()
diff --git a/server/site_server_job_utils.py b/server/site_server_job_utils.py
index 0384f4b..0eb5bf6 100644
--- a/server/site_server_job_utils.py
+++ b/server/site_server_job_utils.py
@@ -204,10 +204,10 @@
         # the machine we're working on. Required so that server jobs will write
         # to the proper location.
         self._server_job.machines = [self._machine]
-        self._server_job.push_execution_context(self._machine)
+        self._server_job.push_execution_context(self._machine['hostname'])
         os.chdir(self._server_job.resultdir)
         if self._continuous_parsing:
-            self._server_job._parse_job += "/" + self._machine
+            self._server_job._parse_job += "/" + self._machine['hostname']
             self._server_job._using_parser = True
             self._server_job.init_parser()
 
diff --git a/server/site_utils.py b/server/site_utils.py
index 2d2d050..1ebff49 100644
--- a/server/site_utils.py
+++ b/server/site_utils.py
@@ -670,4 +670,24 @@
 def verify_not_root_user():
     """Simple function to error out if running with uid == 0"""
     if os.getuid() == 0:
-        raise error.IllegalUser('This script can not be ran as root.')
\ No newline at end of file
+        raise error.IllegalUser('This script can not be ran as root.')
+
+
+def get_hostname_from_machine(machine):
+    """Lookup hostname from a machine string or dict.
+
+    @returns: Machine hostname in string format.
+    """
+    hostname, _ = get_host_info_from_machine(machine)
+    return hostname
+
+
+def get_host_info_from_machine(machine):
+    """Lookup host information from a machine string or dict.
+
+    @returns: Tuple of (hostname, host_attributes)
+    """
+    if isinstance(machine, dict):
+        return (machine['hostname'], machine['host_attributes'])
+    else:
+        return (machine, {})
diff --git a/server/test.py b/server/test.py
index b1e8641..a46b472 100644
--- a/server/test.py
+++ b/server/test.py
@@ -97,7 +97,7 @@
     def _install(self):
         if not self.host:
             from autotest_lib.server import hosts, autotest
-            self.host = hosts.create_host(self.job.machines[0])
+            self.host = hosts.create_host(self.job.machine_dict_list[0])
             # TODO(kevcheng): remove when host client install is supported for
             # ADBHost. crbug.com/543702
             if not self.host.is_client_install_supported: