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')