[autotest] Acquire sonic hosts.

Sonic tests need a way to acquire hosts. Till crbug.com/337887 is completely
solved locking them should serve as a good intermediate solution.

TEST=Ran the sonic tests.
BUG=chromium:337887

Change-Id: I4b844e13565aa95de87f9d978908c0d0c1576a60
Reviewed-on: https://chromium-review.googlesource.com/184968
Tested-by: Prashanth B <beeps@chromium.org>
Reviewed-by: Dan Shi <dshi@chromium.org>
Reviewed-by: Kris Rambish <krisr@chromium.org>
Commit-Queue: Prashanth B <beeps@chromium.org>
diff --git a/server/cros/chaos_lib/chaos_runner.py b/server/cros/chaos_lib/chaos_runner.py
index 742b2b0..8926742 100644
--- a/server/cros/chaos_lib/chaos_runner.py
+++ b/server/cros/chaos_lib/chaos_runner.py
@@ -5,14 +5,13 @@
 import contextlib
 import datetime
 import logging
-import random
 
-from autotest_lib.client.common_lib import error
 from autotest_lib.client.common_lib.cros.network import chaos_constants
 from autotest_lib.client.common_lib.cros.network import iw_runner
 from autotest_lib.server import hosts
 from autotest_lib.server import frontend
 from autotest_lib.server import site_linux_system
+from autotest_lib.server import site_utils
 from autotest_lib.server.cros import host_lock_manager
 from autotest_lib.server.cros.chaos_ap_configurators import ap_batch_locker
 from autotest_lib.server.cros.chaos_ap_configurators import ap_cartridge
@@ -51,28 +50,14 @@
         @param lock_manager HostLockManager object.
         @param hostname string optional hostname of a packet capture machine.
 
+        @return: An SSHHost object representing a locked packet_capture machine.
         """
         if hostname is not None:
             return hosts.SSHHost(hostname)
 
         afe = frontend.AFE(debug=True)
-        potential_hosts = afe.get_hosts(multiple_labels=['packet_capture'])
-        if not potential_hosts:
-            raise error.TestError('No packet capture machines available.')
-
-        # Shuffle hosts so that we don't lock the same packet capture host
-        # every time.  This prevents errors where a fault might seem repeatable
-        # because we lock the same packet capturer for each test run.
-        random.shuffle(potential_hosts)
-        for host in potential_hosts:
-            if lock_manager.lock([host.hostname]):
-                logging.info('Locked packet capture host %s.', host.hostname)
-                return hosts.SSHHost(host.hostname + '.cros')
-            else:
-                logging.info('Unable to lock packet capture host %s.',
-                             host.hostname)
-
-        raise error.TestError('Could not allocate a packet tracer.')
+        return hosts.SSHHost(site_utils.lock_host_with_labels(
+                afe, lock_manager, labels=['packet_capture']) + '.cros')
 
 
     def _power_down_aps(self, aps):
diff --git a/server/cros/sonic_client_utils.py b/server/cros/sonic_client_utils.py
index 72d0403..2d4ef70 100644
--- a/server/cros/sonic_client_utils.py
+++ b/server/cros/sonic_client_utils.py
@@ -21,8 +21,11 @@
 
 import common
 
+from autotest_lib.client.bin import utils
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.common_lib.cros import retry
+from autotest_lib.server import frontend
+from autotest_lib.server import site_utils
 
 
 # Give all our rpcs about six seconds of retry time. If a longer timeout
@@ -36,6 +39,7 @@
                 'cD/djLLQm2zZfFygP4U4/o++ZM91EWtgII10LisoS47qT2TIOg4Un4+G57e'
                 'lZ9PjEIhcJfANqkYrD3t9dpEzMNr936TLB2u683B5qmbB68Nq1Eel7KVc+F'
                 '0BqhBondDqhvDvGPEV0vBsbErJFlNH7SQIDAQAB')
+SONIC_BOARD_LABEL = 'board:sonic'
 
 
 def get_extension_id(pub_key_pem=MANIFEST_KEY):
@@ -131,6 +135,33 @@
     urllib2.urlopen(request)
 
 
+@retry.retry(RPC_EXCEPTIONS + (error.TestError,), timeout_min=30)
+def acquire_sonic(lock_manager, additional_labels=None):
+    """Lock a host that has the sonic host labels.
+
+    @param lock_manager: A manager for locking/unlocking hosts, as defined by
+        server.cros.host_lock_manager.
+    @param additional_labels: A list of additional labels to apply in the search
+        for a sonic device.
+
+    @return: A string specifying the hostname of a locked sonic host.
+
+    @raises ValueError: Is no hosts matching the given labels are found.
+    """
+    sonic_host = None
+    afe = frontend.AFE(debug=True)
+    labels = [SONIC_BOARD_LABEL]
+    if additional_labels:
+        labels += additional_labels
+    sonic_hostname = utils.poll_for_condition(
+            lambda: site_utils.lock_host_with_labels(afe, lock_manager, labels),
+            sleep_interval=60,
+            exception=SonicProxyException('Timed out trying to find a sonic '
+                                          'host with labels %s.' % labels))
+    logging.info('Acquired sonic host returned %s', sonic_hostname)
+    return sonic_hostname
+
+
 class SonicProxyException(Exception):
     """Generic exception raised when a sonic rpc fails."""
     pass
diff --git a/server/cros/sonic_extension_downloader.py b/server/cros/sonic_extension_downloader.py
index 2447802..348fa6d 100644
--- a/server/cros/sonic_extension_downloader.py
+++ b/server/cros/sonic_extension_downloader.py
@@ -5,6 +5,7 @@
 
 import httplib2
 import json
+import logging
 import os
 import re
 import shutil
@@ -44,6 +45,8 @@
     response_xml = httplib2.Http().request(update_check_link, 'GET')[1]
     codebase_match = re.compile(r'codebase="(.*crx)"').search(response_xml)
     if codebase_match is not None:
+        logging.info('Omaha response while downloading extension: %s',
+                     response_xml)
         return codebase_match.groups()[0]
     raise IOError('Omaha response is invalid %s.' % response_xml)
 
@@ -54,6 +57,7 @@
     @param dest_file: Path to a destination file for the extension.
     """
     download_url = get_download_url_from_omaha(TEST_EXTENSION_ID)
+    logging.info('Downloading extension from %s', download_url)
     response = urllib2.urlopen(download_url)
     with open(dest_file, 'w') as f:
         f.write(response.read())
@@ -110,6 +114,8 @@
 
     if not os.path.exists(unzipped_crx_dir):
         raise SonicDownloaderException('Unable to download sonic extension.')
+    logging.info('Sonic extension successfully downloaded into %s.',
+                 unzipped_crx_dir)
 
     # TODO(beeps): crbug.com/325869, investigate the limits of component
     # extensions. For now this is ok because even sonic testing inlines a
diff --git a/server/hosts/sonic_host.py b/server/hosts/sonic_host.py
index 600acd1..e5c906a 100644
--- a/server/hosts/sonic_host.py
+++ b/server/hosts/sonic_host.py
@@ -39,6 +39,9 @@
     OTA_LOCATION = '/cache/ota.zip'
     RECOVERY_DIR = '/cache/recovery'
     COMMAND_FILE = os.path.join(RECOVERY_DIR, 'command')
+    PLATFORM = 'sonic'
+    LABELS = [sonic_client_utils.SONIC_BOARD_LABEL]
+
 
     @staticmethod
     def check_host(host, timeout=10):
@@ -116,6 +119,14 @@
         return self.run(cmd, timeout=timeout).stdout.strip()
 
 
+    def get_platform(self):
+        return self.PLATFORM
+
+
+    def get_labels(self):
+        return self.LABELS
+
+
     def ssh_ping(self, timeout=60, base_cmd=''):
         """Checks if we can ssh into the host and run getprop.
 
diff --git a/server/site_tests/sonic_AppTest/control b/server/site_tests/sonic_AppTest/control
index 8ba7b84..b5f662b 100644
--- a/server/site_tests/sonic_AppTest/control
+++ b/server/site_tests/sonic_AppTest/control
@@ -2,8 +2,13 @@
 # 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 utils
+from autotest_lib.server.cros import host_lock_manager
+from autotest_lib.server.cros import sonic_client_utils
+
 
 AUTHOR = "Chrome OS Team"
 NAME = "sonic_AppTest"
@@ -30,16 +35,25 @@
 args_dict = utils.args_to_dict(args)
 
 def run(machine):
-    sonic_hostname = args_dict.get('sonic_hostname')
     extension_dir = args_dict.get('extension_dir')
-    if not sonic_hostname:
-        raise error.TestError('Cannot run sonic_AppTest without a sonic host. '
-                              'please specify --args="sonic_hostname=<ip>" with '
-                              'test_that.')
+    lock_manager = host_lock_manager.HostLockManager()
 
-    cros_host = hosts.create_host(machine)
-    sonic_host = hosts.create_host(sonic_hostname)
-    job.run_test('sonic_AppTest', cros_host=cros_host, sonic_host=sonic_host,
-                 extension_dir=extension_dir, disable_sysinfo=True)
+    # If the hostname of a sonic device to use with this test
+    # is passed through --args, just try to lock it, otherwise
+    # look for an unlocked host with the sonic label. The context
+    # manager handles unlocking hosts locked through the lock_manager.
+    with host_lock_manager.HostsLockedBy(lock_manager):
+        sonic_hostname = args_dict.get('sonic_hostname')
+        if sonic_hostname:
+            lock_manager.lock([sonic_hostname])
+        else:
+            sonic_hostname = sonic_client_utils.acquire_sonic(lock_manager)
+        logging.info('Using sonic host: %s', sonic_hostname)
+        cros_host = hosts.create_host(machine)
+        sonic_host = hosts.create_host(sonic_hostname)
+        job.run_test('sonic_AppTest', cros_host=cros_host,
+                     sonic_host=sonic_host, extension_dir=extension_dir,
+                     disable_sysinfo=True)
+
 
 parallel_simple(run, machines)
diff --git a/server/site_tests/sonic_AppTest/sonic_AppTest.py b/server/site_tests/sonic_AppTest/sonic_AppTest.py
index cb0184d..3ac3244 100644
--- a/server/site_tests/sonic_AppTest/sonic_AppTest.py
+++ b/server/site_tests/sonic_AppTest/sonic_AppTest.py
@@ -19,14 +19,18 @@
 
 
     def initialize(self, sonic_host, extension_dir=None):
-        """Download the latest extension."""
-        # TODO(beeps): crbug.com/337708
+        """Download the latest extension, or use a local path if specified."""
+        # TODO: crbug.com/337708
         if not extension_dir:
+            logging.info('Downloading ToT extension for test since no local '
+                         'extension specified.')
             extension_path = os.path.join(self.job.clientdir, 'deps',
                                           'sonic_extension')
             sonic_extension_downloader.setup_extension(extension_path)
             self.extension_path = extension_path
         else:
+            logging.info('Using local extension for test %s.',
+                         self.extension_dir)
             self.extension_path = None
 
 
@@ -49,6 +53,8 @@
         @raises TestError: If the app didn't start, or the app was unrecognized,
             or the payload is invalid.
         """
+        logging.info('Testing app %s, sonic_host %s and chromeos device %s ',
+                     app, sonic_host.hostname, cros_host.hostname)
         if app == 'ChromeCast':
             sonic_host.enable_test_extension()
             client_at = autotest.Autotest(cros_host)
diff --git a/server/site_utils.py b/server/site_utils.py
index d27aff3..55b1cc6 100644
--- a/server/site_utils.py
+++ b/server/site_utils.py
@@ -6,12 +6,15 @@
 import httplib
 import json
 import logging
+import random
 import re
 import time
 import urllib2
 
 import common
-from autotest_lib.client.common_lib import base_utils, global_config
+from autotest_lib.client.common_lib import base_utils
+from autotest_lib.client.common_lib import error
+from autotest_lib.client.common_lib import global_config
 from autotest_lib.server.cros.dynamic_suite import constants
 
 
@@ -232,3 +235,40 @@
         logging.warn('Could not get a status from %s', status_url)
         return
     _decode_lab_status(json_status, build)
+
+
+def lock_host_with_labels(afe, lock_manager, labels):
+    """Lookup and lock one host that matches the list of input labels.
+
+    @param afe: An instance of the afe class, as defined in server.frontend.
+    @param lock_manager: A lock manager capable of locking hosts, eg the
+        one defined in server.cros.host_lock_manager.
+    @param labels: A list of labels to look for on hosts.
+
+    @return: The hostname of a host matching all labels, and locked through the
+        lock_manager. The hostname will be as specified in the database the afe
+        object is associated with, i.e if it exists in afe_hosts with a .cros
+        suffix, the hostname returned will contain a .cros suffix.
+
+    @raises: error.NoEligibleHostException: If no hosts matching the list of
+        input labels are available.
+    @raises: error.TestError: If unable to lock a host matching the labels.
+    """
+    potential_hosts = afe.get_hosts(multiple_labels=labels)
+    if not potential_hosts:
+        raise error.NoEligibleHostException(
+                'No devices found with labels %s.' % labels)
+
+    # This prevents errors where a fault might seem repeatable
+    # because we lock, say, the same packet capturer for each test run.
+    random.shuffle(potential_hosts)
+    for host in potential_hosts:
+        if lock_manager.lock([host.hostname]):
+            logging.info('Locked device %s with labels %s.',
+                         host.hostname, labels)
+            return host.hostname
+        else:
+            logging.info('Unable to lock device %s with labels %s.',
+                         host.hostname, labels)
+
+    raise error.TestError('Could not lock a device with labels %s' % labels)