brillo: Connect to wifi before running wifi tests.

The brillo_PingTest and brillo_WifiInterfaceTest require wifi to be
connected before running in the lab.

BUG:b:27297360
TEST=tests pass when ssid arg is GoogleGuest

Change-Id: Iff39c5b182f4b94a2bc5fac04bb3bf7710324321
Reviewed-on: https://chromium-review.googlesource.com/334404
Commit-Ready: Ralph Nathan <ralphnathan@chromium.org>
Tested-by: Ralph Nathan <ralphnathan@chromium.org>
Reviewed-by: Ralph Nathan <ralphnathan@chromium.org>
diff --git a/server/brillo/host_utils.py b/server/brillo/host_utils.py
index 7ac30bf..05a3b02 100644
--- a/server/brillo/host_utils.py
+++ b/server/brillo/host_utils.py
@@ -4,6 +4,14 @@
 
 """Utilities used with Brillo hosts."""
 
+import contextlib
+import logging
+import time
+
+import common
+from autotest_lib.client.bin import site_utils as client_site_utils
+
+
 _RUN_BACKGROUND_TEMPLATE = '( %(cmd)s ) </dev/null >/dev/null 2>&1 & echo -n $!'
 
 _WAIT_CMD_TEMPLATE = """\
@@ -40,3 +48,43 @@
     """
     wait_cmd = _WAIT_CMD_TEMPLATE % {'pid': pid, 'timeout': timeout}
     return host.run(wait_cmd, ignore_status=True).exit_status == 0
+
+
+@contextlib.contextmanager
+def connect_to_ssid(host, ssid, passphrase):
+    """Connects to a given ssid.
+
+    @param host: A host object representing the DUT.
+    @param ssid: A string representing the ssid to connect to.
+    @param passphrase: A string representing the passphrase to the ssid.
+                       Defaults to None.
+    """
+    try:
+        # Update the weaved init.rc to stop privet. This sets up shill in client
+        # mode allowing it to connect to wifi.
+        host.remount()
+        host.run('sed \'s/service weaved \/system\/bin\/weaved/'
+                 'service weaved \/system\/bin\/weaved --disable_privet/\' -i '
+                 '/system/etc/init/weaved.rc')
+        host.reboot()
+        client_site_utils.poll_for_condition(
+                lambda: 'running' in host.run('getprop init.svc.shill').stdout,
+                sleep_interval=1, timeout=300,
+                desc='shill was not started by init')
+        logging.info('Connecting to wifi')
+        wifi_cmd = 'su root shill_setup_wifi --ssid=%s' % ssid
+        if passphrase:
+            wifi_cmd += ' --passphrase=%s' % passphrase
+        host.run(wifi_cmd)
+        # TODO(ralphnathan): Once shill_setup_wifi can monitor the service as it
+        # connects to wifi, remove this timeout.
+        # Wait for wifi connection to occur.
+        time.sleep(10)
+        yield
+    finally:
+        host.remount()
+        host.run('sed \'s/service weaved \/system\/bin\/weaved '
+                 '--disable_privet/service weaved \/system\/bin\/weaved/\' -i '
+                 '/system/etc/init/weaved.rc')
+        host.run('su root stop weaved')
+        host.run('su root start weaved')
diff --git a/server/site_tests/brillo_PingTest/brillo_PingTest.py b/server/site_tests/brillo_PingTest/brillo_PingTest.py
index 7897c80..d7f2d34 100644
--- a/server/site_tests/brillo_PingTest/brillo_PingTest.py
+++ b/server/site_tests/brillo_PingTest/brillo_PingTest.py
@@ -2,9 +2,13 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import logging
+
 import common
 from autotest_lib.client.common_lib import error
+from autotest_lib.client.common_lib import site_utils
 from autotest_lib.server import test
+from autotest_lib.server.brillo import host_utils
 
 
 _DEFAULT_PING_HOST = 'www.google.com'
@@ -16,12 +20,15 @@
     """Ping an Internet host."""
     version = 1
 
-    def run_once(self, host=None, ping_host=_DEFAULT_PING_HOST,
+    def run_once(self, host=None, ssid=None, passphrase=None,
+                 ping_host=_DEFAULT_PING_HOST,
                  ping_count=_DEFAULT_PING_COUNT,
                  ping_timeout=_DEFAULT_PING_TIMEOUT):
         """Pings an Internet host with given timeout and count values.
 
         @param host: A host object representing the DUT.
+        @param ssid: Ssid to connect to.
+        @param passphrase: A string representing the passphrase to the ssid.
         @param ping_host: The Internet host to ping.
         @param ping_count: The number of pings to attempt. The test passes if
                            we get at least one reply.
@@ -29,10 +36,15 @@
 
         @raise TestFail: The test failed.
         """
-        cmd = 'ping -q -c %s -W %s %s' % (ping_count, ping_timeout, ping_host)
-        try:
-            host.run(cmd)
-        except error.AutoservRunError:
-            raise error.TestFail(
-                    'Failed to ping %s in %d seconds on all %d attempts' %
-                    (ping_host, ping_timeout, ping_count))
+        if ssid is None:
+            ssid = site_utils.get_wireless_ssid(host.hostname)
+        logging.info('Connecting to ssid %s', ssid)
+        with host_utils.connect_to_ssid(host, ssid, passphrase):
+            cmd = 'ping -q -c %s -W %s %s' % (ping_count, ping_timeout,
+                                              ping_host)
+            try:
+                host.run(cmd)
+            except error.AutoservRunError:
+                raise error.TestFail(
+                        'Failed to ping %s in %d seconds on all %d attempts' %
+                        (ping_host, ping_timeout, ping_count))
diff --git a/server/site_tests/brillo_PingTest/control b/server/site_tests/brillo_PingTest/control
index 1d23a25..9a304ee 100644
--- a/server/site_tests/brillo_PingTest/control
+++ b/server/site_tests/brillo_PingTest/control
@@ -15,7 +15,8 @@
 Tests whether a Brillo device can ping the Internet.
 """
 
-TEST_ARG_NAMES = ('ping_host', 'ping_count', 'ping_timeout')
+TEST_ARG_NAMES = ('ssid', 'passphrase', 'ping_host', 'ping_count',
+                  'ping_timeout')
 args_dict = utils.args_to_dict(args)
 
 def run(machine):
diff --git a/server/site_tests/brillo_WifiInterfaceTest/brillo_WifiInterfaceTest.py b/server/site_tests/brillo_WifiInterfaceTest/brillo_WifiInterfaceTest.py
index 5865088..1327d75 100644
--- a/server/site_tests/brillo_WifiInterfaceTest/brillo_WifiInterfaceTest.py
+++ b/server/site_tests/brillo_WifiInterfaceTest/brillo_WifiInterfaceTest.py
@@ -6,8 +6,10 @@
 
 import common
 from autotest_lib.client.common_lib import error
+from autotest_lib.client.common_lib import site_utils
 from autotest_lib.client.common_lib.cros.network import iw_runner
 from autotest_lib.server import test
+from autotest_lib.server.brillo import host_utils
 
 
 class brillo_WifiInterfaceTest(test.test):
@@ -38,10 +40,13 @@
         return ifconfig_dict
 
 
-    def run_once(self, host=None, wifi_iface=None, wifi_ssid=None):
+    def run_once(self, host=None, ssid=None, passphrase=None, wifi_iface=None,
+                 wifi_ssid=None):
         """Check that a given wifi interface is properly configured.
 
         @param host: a host object representing the DUT.
+        @param ssid: A string representing the ssid to connect to.
+        @param passphrase: A string representing the passphrase to the ssid.
         @param wifi_iface: Name of the wifi interface to test; None means we'll
                            try to detect at least one that works.
         @param wifi_ssid: Name of the SSID we want the interface to be
@@ -49,49 +54,57 @@
 
         @raise TestFail: The test failed.
         """
-        err_iface = ('No interface is' if wifi_iface is None
-                      else 'Interface %s is not' % wifi_iface)
+        self.host = host
+        if ssid is None:
+            ssid = site_utils.get_wireless_ssid(host.hostname)
+        logging.info('Connecting to ssid %s', ssid)
+        with host_utils.connect_to_ssid(host, ssid, passphrase):
+            err_iface = ('No interface is' if wifi_iface is None
+                          else 'Interface %s is not' % wifi_iface)
 
-        # First check link status and SSID.
-        iw = iw_runner.IwRunner(remote_host=host)
-        active_ifaces = []
-        try:
-            iw_ifaces = [iface_tuple.if_name
-                         for iface_tuple in iw.list_interfaces()]
-            if wifi_iface is not None:
-                if wifi_iface not in iw_ifaces:
-                    raise error.TestFail(
-                            'Interface %s not listed by iw' % wifi_iface)
-                test_ifaces = [wifi_iface]
-            else:
-                test_ifaces = iw_ifaces
+            # First check link status and SSID.
+            iw = iw_runner.IwRunner(remote_host=host)
+            active_ifaces = []
+            try:
+                iw_ifaces = [iface_tuple.if_name
+                             for iface_tuple in iw.list_interfaces()]
+                if wifi_iface is not None:
+                    if wifi_iface not in iw_ifaces:
+                        raise error.TestFail(
+                                'Interface %s not listed by iw' % wifi_iface)
+                    test_ifaces = [wifi_iface]
+                else:
+                    test_ifaces = iw_ifaces
 
-            for iface in test_ifaces:
-                iface_ssid = iw.get_link_value(iface, 'SSID')
-                if (iface_ssid is not None and
-                    (wifi_ssid is None or iface_ssid == wifi_ssid)):
-                    active_ifaces.append(iface)
-        except error.AutoservRunError:
-            raise error.TestFail('Failed to run iw')
+                for iface in test_ifaces:
+                    iface_ssid = iw.get_link_value(iface, 'SSID')
+                    if (iface_ssid is not None and
+                        (wifi_ssid is None or iface_ssid == wifi_ssid)):
+                        active_ifaces.append(iface)
+            except error.AutoservRunError:
+                raise error.TestFail('Failed to run iw')
 
-        if not active_ifaces:
-            err_ssid = 'any SSID' if wifi_ssid is None else 'SSID ' + wifi_ssid
-            raise error.TestFail('%s connected to %s' % (err_iface, err_ssid))
+            if not active_ifaces:
+                err_ssid = 'any SSID' if wifi_ssid is None else ('SSID ' +
+                                                                 wifi_ssid)
+                raise error.TestFail('%s connected to %s' % (err_iface,
+                                                             err_ssid))
 
-        logging.info('Active wifi interfaces: %s', ', '.join(active_ifaces))
+            logging.info('Active wifi interfaces: %s', ', '.join(active_ifaces))
 
-        # Then check IPv4 connectivity.
-        try:
-            ifconfig_output = host.run_output('ifconfig').splitlines()
-        except error.AutoservRunError:
-            raise error.TestFail('Failed to run ifconfig')
+            # Then check IPv4 connectivity.
+            try:
+                ifconfig_output = host.run_output('ifconfig').splitlines()
+            except error.AutoservRunError:
+                raise error.TestFail('Failed to run ifconfig')
 
-        ifconfig_dict = self.get_ifconfig_dict(ifconfig_output)
-        connected_ifaces = [iface for iface in active_ifaces
-                            if any(['inet addr:' in line
-                                    for line in ifconfig_dict.get(iface, [])])]
-        if not connected_ifaces:
-            raise error.TestFail('%s IPv4 connected' % err_iface)
+            ifconfig_dict = self.get_ifconfig_dict(ifconfig_output)
+            connected_ifaces = [iface for iface in active_ifaces
+                                if any(['inet addr:' in line
+                                        for line in ifconfig_dict.get(iface,
+                                                                      [])])]
+            if not connected_ifaces:
+                raise error.TestFail('%s IPv4 connected' % err_iface)
 
-        logging.info('IPv4 connected wifi interfaces: %s',
-                     ', '.join(connected_ifaces))
+            logging.info('IPv4 connected wifi interfaces: %s',
+                         ', '.join(connected_ifaces))
diff --git a/server/site_tests/brillo_WifiInterfaceTest/control b/server/site_tests/brillo_WifiInterfaceTest/control
index 458de34..22a2167 100644
--- a/server/site_tests/brillo_WifiInterfaceTest/control
+++ b/server/site_tests/brillo_WifiInterfaceTest/control
@@ -15,7 +15,7 @@
 Tests whether a Brillo device wifi interface is properly configured.
 """
 
-TEST_ARG_NAMES = ('wifi_iface', 'wifi_ssid')
+TEST_ARG_NAMES = ('ssid', 'passphrase', 'wifi_iface', 'wifi_ssid')
 args_dict = utils.args_to_dict(args)
 
 def run(machine):