FAFT: Move wait_for_client and wait_for_client_offline to ModeSwitcher

These two methods are tightly related to ModeSwitcher. This change
helps it to be derived to different implementations.

BUG=chrome-os-partner:39744
TEST=Ran some of the effected FAFT tests.

Change-Id: Id9be1decec3ebc5a843a30b3070a2352546423eb
Reviewed-on: https://chromium-review.googlesource.com/294773
Commit-Ready: Wai-Hong Tam <waihong@chromium.org>
Tested-by: Wai-Hong Tam <waihong@chromium.org>
Reviewed-by: danny chan <dchan@chromium.org>
diff --git a/server/cros/faft/firmware_test.py b/server/cros/faft/firmware_test.py
index 7d660d8..29dbe4b 100644
--- a/server/cros/faft/firmware_test.py
+++ b/server/cros/faft/firmware_test.py
@@ -21,9 +21,7 @@
 from autotest_lib.server.cros.servo import chrome_ec
 
 
-class ConnectionError(Exception):
-    """Raised on an error of connecting DUT."""
-    pass
+ConnectionError = mode_switcher.ConnectionError
 
 
 class FAFTBase(test.test):
@@ -169,49 +167,6 @@
         super(FirmwareTest, self).cleanup()
         logging.info('FirmwareTest cleanup done (id=%s)', self.run_id)
 
-    def wait_for_client(self, timeout=180):
-        """Wait for the client to come back online.
-
-        New remote processes will be launched if their used flags are enabled.
-
-        @param timeout: Time in seconds to wait for the client SSH daemon to
-                        come up.
-        @raise ConnectionError: Failed to connect DUT.
-        """
-        logging.info("-[FAFT]-[ start wait_for_client ]---")
-        # Wait for the system to respond to ping before attempting ssh
-        if not self._client.ping_wait_up(timeout):
-            logging.warning("-[FAFT]-[ system did not respond to ping ]")
-        if self._client.wait_up(timeout):
-            # Check the FAFT client is avaiable.
-            self.faft_client.system.is_available()
-            # Stop update-engine as it may change firmware/kernel.
-            self._stop_service('update-engine')
-        else:
-            logging.error('wait_for_client() timed out.')
-            raise ConnectionError()
-        logging.info("-[FAFT]-[ end wait_for_client ]-----")
-
-    def wait_for_client_offline(self, timeout=60, orig_boot_id=None):
-        """Wait for the client to come offline.
-
-        @param timeout: Time in seconds to wait the client to come offline.
-        @param orig_boot_id: A string containing the original boot id.
-        @raise ConnectionError: Failed to wait DUT offline.
-        """
-        # When running against panther, we see that sometimes
-        # ping_wait_down() does not work correctly. There needs to
-        # be some investigation to the root cause.
-        # If we sleep for 120s before running get_boot_id(), it
-        # does succeed. But if we change this to ping_wait_down()
-        # there are implications on the wait time when running
-        # commands at the fw screens.
-        if not self._client.ping_wait_down(timeout):
-            if orig_boot_id and self._client.get_boot_id() != orig_boot_id:
-                logging.warn('Reboot done very quickly.')
-                return
-            raise ConnectionError()
-
     def _record_system_info(self):
         """Record some critical system info to the attr keyval.
 
@@ -245,7 +200,7 @@
             self.servo.switch_usbkey('dut')
 
         try:
-            self.wait_for_client()
+            self.switcher.wait_for_client()
             lines = self.faft_client.system.run_shell_command_get_output(
                         'crossystem recovery_reason')
             recovery_reason = int(lines[0])
@@ -269,10 +224,10 @@
         self.switcher.mode_aware_reboot(reboot_type='cold',
                                         sync_before_boot=False,
                                         wait_for_dut_up=False)
-        self.wait_for_client_offline()
+        self.switcher.wait_for_client_offline()
         self.switcher.bypass_dev_mode()
         try:
-            self.wait_for_client()
+            self.switcher.wait_for_client()
             return
         except ConnectionError:
             logging.warn('Cold reboot doesn\'t help, still connection error.')
@@ -309,10 +264,10 @@
         logging.info('Try restore the OS image...')
         self.faft_client.system.run_shell_command('chromeos-install --yes')
         self.switcher.mode_aware_reboot(wait_for_dut_up=False)
-        self.wait_for_client_offline()
+        self.switcher.wait_for_client_offline()
         self.switcher.bypass_dev_mode()
         try:
-            self.wait_for_client()
+            self.switcher.wait_for_client()
             logging.info('Successfully restore OS image.')
             return
         except ConnectionError:
@@ -330,7 +285,7 @@
         self.servo.switch_usbkey('host')
         self.switcher.bypass_rec_mode()
         try:
-            self.wait_for_client()
+            self.switcher.wait_for_client()
         except ConnectionError:
             raise error.TestError('Failed to boot the USB image.')
 
@@ -931,7 +886,7 @@
         self.servo.power_key(self.faft_config.hold_pwr_button_poweroff)
         # device can take 44-51 seconds to restart,
         # add buffer from the default timeout of 60 seconds.
-        self.wait_for_client_offline(timeout=100, orig_boot_id=boot_id)
+        self.switcher.wait_for_client_offline(timeout=100, orig_boot_id=boot_id)
         time.sleep(self.faft_config.shutdown)
         # Short press power button to boot DUT again.
         self.servo.power_key(self.faft_config.hold_pwr_button_poweron)
@@ -1059,7 +1014,7 @@
         try:
             if shutdown_timeout is None:
                 shutdown_timeout = self.faft_config.shutdown_timeout
-            self.wait_for_client(timeout=shutdown_timeout)
+            self.switcher.wait_for_client(timeout=shutdown_timeout)
             raise error.TestFail(
                     'Should shut the device down after calling %s.' %
                     str(shutdown_action))
diff --git a/server/cros/faft/utils/mode_switcher.py b/server/cros/faft/utils/mode_switcher.py
index c6d07d2..125648a 100644
--- a/server/cros/faft/utils/mode_switcher.py
+++ b/server/cros/faft/utils/mode_switcher.py
@@ -6,6 +6,11 @@
 import time
 
 
+class ConnectionError(Exception):
+    """Raised on an error of connecting DUT."""
+    pass
+
+
 class _BaseFwBypasser(object):
     """Base class that controls bypass logic for firmware screens."""
 
@@ -172,6 +177,7 @@
 
     def __init__(self, faft_framework):
         self.faft_framework = faft_framework
+        self.client_host = faft_framework._client
         self.faft_client = faft_framework.faft_client
         self.servo = faft_framework.servo
         self.faft_config = faft_framework.faft_config
@@ -224,18 +230,18 @@
             self._enable_rec_mode_and_reboot(usb_state='dut')
             if wait_for_dut_up:
                 self.bypasser.bypass_rec_mode()
-                self.faft_framework.wait_for_client()
+                self.wait_for_client()
 
         elif to_mode == 'dev':
             self._enable_dev_mode_and_reboot()
             if wait_for_dut_up:
                 self.bypasser.bypass_dev_mode()
-                self.faft_framework.wait_for_client()
+                self.wait_for_client()
 
         elif to_mode == 'normal':
             self._enable_normal_mode_and_reboot()
             if wait_for_dut_up:
-                self.faft_framework.wait_for_client()
+                self.wait_for_client()
 
         else:
             raise NotImplementedError(
@@ -279,7 +285,7 @@
             self.faft_framework.blocking_sync()
         reboot_method()
         if sync_before_boot:
-            self.faft_framework.wait_for_client_offline(orig_boot_id=boot_id)
+            self.wait_for_client_offline(orig_boot_id=boot_id)
         if wait_for_dut_up:
             # For encapsulating the behavior of skipping firmware screen,
             # e.g. requiring unplug and plug USB, the variants are not
@@ -295,7 +301,7 @@
                 self.bypasser.bypass_dev_mode()
             if not is_dev:
                 self.bypasser.bypass_rec_mode()
-            self.faft_framework.wait_for_client()
+            self.wait_for_client()
         logging.info("-[ModeSwitcher]-[ end mode_aware_reboot(%r, %s, ..) ]-",
                      reboot_type, reboot_method.__name__)
 
@@ -368,6 +374,51 @@
         self.bypasser.trigger_dev_to_normal()
 
 
+    def wait_for_client(self, timeout=180):
+        """Wait for the client to come back online.
+
+        New remote processes will be launched if their used flags are enabled.
+
+        @param timeout: Time in seconds to wait for the client SSH daemon to
+                        come up.
+        @raise ConnectionError: Failed to connect DUT.
+        """
+        logging.info("-[FAFT]-[ start wait_for_client ]---")
+        # Wait for the system to respond to ping before attempting ssh
+        if not self.client_host.ping_wait_up(timeout):
+            logging.warning("-[FAFT]-[ system did not respond to ping ]")
+        if self.client_host.wait_up(timeout):
+            # Check the FAFT client is avaiable.
+            self.faft_client.system.is_available()
+            # Stop update-engine as it may change firmware/kernel.
+            self.faft_framework._stop_service('update-engine')
+        else:
+            logging.error('wait_for_client() timed out.')
+            raise ConnectionError()
+        logging.info("-[FAFT]-[ end wait_for_client ]-----")
+
+
+    def wait_for_client_offline(self, timeout=60, orig_boot_id=None):
+        """Wait for the client to come offline.
+
+        @param timeout: Time in seconds to wait the client to come offline.
+        @param orig_boot_id: A string containing the original boot id.
+        @raise ConnectionError: Failed to wait DUT offline.
+        """
+        # When running against panther, we see that sometimes
+        # ping_wait_down() does not work correctly. There needs to
+        # be some investigation to the root cause.
+        # If we sleep for 120s before running get_boot_id(), it
+        # does succeed. But if we change this to ping_wait_down()
+        # there are implications on the wait time when running
+        # commands at the fw screens.
+        if not self.client_host.ping_wait_down(timeout):
+            if orig_boot_id and self.client_host.get_boot_id() != orig_boot_id:
+                logging.warn('Reboot done very quickly.')
+                return
+            raise ConnectionError()
+
+
 class _PhysicalButtonSwitcher(_BaseModeSwitcher):
     """Class that switches firmware mode via physical button."""
 
@@ -394,7 +445,7 @@
         # Rebooting EC with rec mode on. Should power on AP.
         # Plug out USB disk for preventing recovery boot without warning
         self._enable_rec_mode_and_reboot(usb_state='host')
-        self.faft_framework.wait_for_client_offline()
+        self.wait_for_client_offline()
         self.bypasser.trigger_rec_to_dev()
 
 
@@ -402,7 +453,7 @@
         """Switch to normal mode and reboot."""
         logging.info("Disabling keyboard controlled developer mode")
         self._disable_rec_mode_and_reboot()
-        self.faft_framework.wait_for_client_offline()
+        self.wait_for_client_offline()
         self.bypasser.trigger_dev_to_normal()
 
 
@@ -413,7 +464,7 @@
         """Switch to developer mode and reboot."""
         logging.info("Enabling Jetstream developer mode")
         self._enable_rec_mode_and_reboot(usb_state='host')
-        self.faft_framework.wait_for_client_offline()
+        self.wait_for_client_offline()
         self.bypasser.trigger_rec_to_dev()
 
 
diff --git a/server/site_tests/firmware_ConsecutiveBoot/firmware_ConsecutiveBoot.py b/server/site_tests/firmware_ConsecutiveBoot/firmware_ConsecutiveBoot.py
index 1ba4bf5..34774e3 100644
--- a/server/site_tests/firmware_ConsecutiveBoot/firmware_ConsecutiveBoot.py
+++ b/server/site_tests/firmware_ConsecutiveBoot/firmware_ConsecutiveBoot.py
@@ -45,9 +45,9 @@
         """
         logging.info('wait_for_client %d start.', self.faft_waitup_time)
         if self.faft_waitup_time:
-            self.wait_for_client(self.faft_waitup_time)
+            self.switcher.wait_for_client(self.faft_waitup_time)
         else:
-            self.wait_for_client()
+            self.switcher.wait_for_client()
 
     def shutdown_power_on(self):
         """
@@ -61,7 +61,7 @@
         logging.info("Sending /sbin/shutdown -P now")
         self.faft_client.system.run_shell_command('/sbin/shutdown -P now')
         logging.info('Wait for client to go offline')
-        self.wait_for_client_offline(timeout=100, orig_boot_id=boot_id)
+        self.switcher.wait_for_client_offline(timeout=100, orig_boot_id=boot_id)
         # Retry in case power_short_press was not registered.
         for i in xrange(self.POWER_ON_RETRY):
             logging.info("sleep %d, tap power key to boot.",
diff --git a/server/site_tests/firmware_ConsecutiveBootPowerButton/firmware_ConsecutiveBootPowerButton.py b/server/site_tests/firmware_ConsecutiveBootPowerButton/firmware_ConsecutiveBootPowerButton.py
index d70ea95..46683d6 100644
--- a/server/site_tests/firmware_ConsecutiveBootPowerButton/firmware_ConsecutiveBootPowerButton.py
+++ b/server/site_tests/firmware_ConsecutiveBootPowerButton/firmware_ConsecutiveBootPowerButton.py
@@ -39,7 +39,7 @@
                         'mainfw_type': 'developer' if dev_mode else 'normal',
                         }))
             self.full_power_off_and_on()
-            self.wait_for_client()
+            self.switcher.wait_for_client()
 
             logging.info("Expected boot fine.")
             self.check_state((self.checkers.crossystem_checker, {
diff --git a/server/site_tests/firmware_DevBootUSB/firmware_DevBootUSB.py b/server/site_tests/firmware_DevBootUSB/firmware_DevBootUSB.py
index ef0ef24..71cdfb7 100644
--- a/server/site_tests/firmware_DevBootUSB/firmware_DevBootUSB.py
+++ b/server/site_tests/firmware_DevBootUSB/firmware_DevBootUSB.py
@@ -3,7 +3,6 @@
 # found in the LICENSE file.
 
 import logging
-import time
 
 from autotest_lib.client.common_lib import error
 from autotest_lib.server.cros.faft.firmware_test import FirmwareTest
@@ -64,7 +63,7 @@
         self.faft_client.system.set_dev_boot_usb(1)
         self.switcher.mode_aware_reboot(wait_for_dut_up=False)
         self.switcher.bypass_dev_boot_usb()
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Expected USB boot, set dev_boot_usb to the original.")
         self.check_state((self.checkers.dev_boot_usb_checker,
@@ -73,7 +72,7 @@
         self.faft_client.system.set_dev_boot_usb(self.original_dev_boot_usb)
         self.switcher.mode_aware_reboot(wait_for_dut_up=False)
         self.switcher.bypass_dev_mode()
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Check and done.")
         self.check_state((self.checkers.dev_boot_usb_checker, False))
diff --git a/server/site_tests/firmware_DevScreenTimeout/firmware_DevScreenTimeout.py b/server/site_tests/firmware_DevScreenTimeout/firmware_DevScreenTimeout.py
index 7a41e11..a78c4db 100644
--- a/server/site_tests/firmware_DevScreenTimeout/firmware_DevScreenTimeout.py
+++ b/server/site_tests/firmware_DevScreenTimeout/firmware_DevScreenTimeout.py
@@ -100,14 +100,14 @@
         logging.info("Reboot and press Ctrl-D repeatedly.")
         self.switcher.mode_aware_reboot(wait_for_dut_up=False)
         self.ctrl_d_repeatedly()
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Record the firmware boot time without waiting "
                      "firmware screen; on next reboot, do nothing and wait the "
                      "screen timeout.")
         self.record_fw_boot_time('ctrl_d_boot')
         self.switcher.mode_aware_reboot(wait_for_dut_up=False)
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Check the firmware screen timeout matches our spec.")
         self.check_timeout_period()
diff --git a/server/site_tests/firmware_DevTriggerRecovery/firmware_DevTriggerRecovery.py b/server/site_tests/firmware_DevTriggerRecovery/firmware_DevTriggerRecovery.py
index b587779..ee79529 100644
--- a/server/site_tests/firmware_DevTriggerRecovery/firmware_DevTriggerRecovery.py
+++ b/server/site_tests/firmware_DevTriggerRecovery/firmware_DevTriggerRecovery.py
@@ -73,7 +73,7 @@
         self.servo.enable_development_mode()
         self.switcher.mode_aware_reboot(wait_for_dut_up=False)
         self.switcher.bypass_dev_mode()
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Expected values based on platforms (see above), "
                      "run 'chromeos-firmwareupdate --mode todev && reboot', "
@@ -84,7 +84,7 @@
         # Ignore the default reboot_action here because the
         # userspace_action (firmware updater) will reboot the system.
         self.switcher.trigger_dev_to_rec()
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Expected recovery boot and disable dev switch.")
         self.check_state((self.checkers.crossystem_checker, {
@@ -101,7 +101,7 @@
         self.check_state(self.check_devsw_off_transition)
         self.faft_client.system.run_shell_command(
                             'chromeos-firmwareupdate --mode tonormal && reboot')
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Expected normal mode boot, done.")
         self.check_state((self.checkers.crossystem_checker, {
diff --git a/server/site_tests/firmware_ECUsbPorts/firmware_ECUsbPorts.py b/server/site_tests/firmware_ECUsbPorts/firmware_ECUsbPorts.py
index ebdb840..9710153 100644
--- a/server/site_tests/firmware_ECUsbPorts/firmware_ECUsbPorts.py
+++ b/server/site_tests/firmware_ECUsbPorts/firmware_ECUsbPorts.py
@@ -119,7 +119,7 @@
         """Shutdown the system and check USB ports are disabled."""
         self._failed = False
         self.faft_client.system.run_shell_command("shutdown -P now")
-        self.wait_for_client_offline()
+        self.switcher.wait_for_client_offline()
         if not self.wait_port_disabled(self._port_count, self.SHUTDOWN_TIMEOUT):
             logging.info("Fails to wait for USB port disabled")
             self._failed = True
diff --git a/server/site_tests/firmware_EventLog/firmware_EventLog.py b/server/site_tests/firmware_EventLog/firmware_EventLog.py
index bd0004c..aea637a 100644
--- a/server/site_tests/firmware_EventLog/firmware_EventLog.py
+++ b/server/site_tests/firmware_EventLog/firmware_EventLog.py
@@ -89,7 +89,7 @@
         logging.info('Verifying eventlog behavior in recovery mode')
         self._cutoff_time = self._now()
         self.switcher.reboot_to_mode(to_mode='rec', wait_for_dut_up=False)
-        self.wait_for_client_offline()
+        self.switcher.wait_for_client_offline()
         time.sleep(self.faft_config.firmware_screen)
         self.switcher.mode_aware_reboot(sync_before_boot=False)
         self.check_state((self.checkers.crossystem_checker, {
diff --git a/server/site_tests/firmware_FwScreenCloseLid/firmware_FwScreenCloseLid.py b/server/site_tests/firmware_FwScreenCloseLid/firmware_FwScreenCloseLid.py
index db41ddb..c6c1806 100644
--- a/server/site_tests/firmware_FwScreenCloseLid/firmware_FwScreenCloseLid.py
+++ b/server/site_tests/firmware_FwScreenCloseLid/firmware_FwScreenCloseLid.py
@@ -88,7 +88,7 @@
         self.run_shutdown_process(self.wait_fw_screen_and_close_lid,
                                   self.servo.lid_open,
                                   self.switcher.bypass_dev_mode)
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Reboot. When the developer screen shown, press "
                      "enter key to trigger either TO_NORM screen (new) or "
@@ -103,7 +103,7 @@
                                   self.servo.lid_open,
                                   self.switcher.bypass_dev_mode,
                                   self.SHORT_SHUTDOWN_CONFIRMATION_PERIOD)
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Request recovery boot. When the RECOVERY INSERT "
                      "screen shows, close lid to make DUT shutdown.")
@@ -117,7 +117,7 @@
                                   self.servo.lid_open,
                                   self.switcher.bypass_dev_mode,
                                   self.SHORT_SHUTDOWN_CONFIRMATION_PERIOD)
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Request recovery boot again. When the recovery "
                      "insert screen shows, insert a corrupted USB and trigger "
@@ -132,7 +132,7 @@
                                   self.servo.lid_open,
                                   self.switcher.bypass_dev_mode,
                                   self.SHORT_SHUTDOWN_CONFIRMATION_PERIOD)
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Switch back to normal mode.")
         self.check_state((self.checkers.crossystem_checker, {
@@ -154,7 +154,7 @@
                                   self.servo.lid_open,
                                   None,
                                   self.SHORT_SHUTDOWN_CONFIRMATION_PERIOD)
-        self.wait_for_client()
+        self.switcher.wait_for_client()
         self.check_state((self.checkers.crossystem_checker, {
                           'devsw_boot': '0',
                           'mainfw_type': 'normal',
diff --git a/server/site_tests/firmware_FwScreenPressPower/firmware_FwScreenPressPower.py b/server/site_tests/firmware_FwScreenPressPower/firmware_FwScreenPressPower.py
index e78f5f7..c33f20a 100644
--- a/server/site_tests/firmware_FwScreenPressPower/firmware_FwScreenPressPower.py
+++ b/server/site_tests/firmware_FwScreenPressPower/firmware_FwScreenPressPower.py
@@ -80,7 +80,7 @@
         self.run_shutdown_process(self.wait_fw_screen_and_press_power,
                                   None,
                                   self.switcher.bypass_dev_mode)
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Reboot. When the developer screen shown, press "
                      "enter key to trigger either TO_NORM screen (new) or "
@@ -95,7 +95,7 @@
                                   None,
                                   self.switcher.bypass_dev_mode,
                                   self.SHORT_SHUTDOWN_CONFIRMATION_PERIOD)
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Request recovery boot. When the RECOVERY INSERT "
                      "screen shows, press power button to make DUT shutdown.")
@@ -109,7 +109,7 @@
                                   None,
                                   self.switcher.bypass_dev_mode,
                                   self.SHORT_SHUTDOWN_CONFIRMATION_PERIOD)
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Request recovery boot again. When the recovery "
                      "insert screen shows, insert a corrupted USB and trigger "
@@ -125,7 +125,7 @@
                                   None,
                                   self.switcher.bypass_dev_mode,
                                   self.SHORT_SHUTDOWN_CONFIRMATION_PERIOD)
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Switch back to normal mode.")
         self.check_state((self.checkers.crossystem_checker, {
@@ -147,7 +147,7 @@
                                   None,
                                   None,
                                   self.SHORT_SHUTDOWN_CONFIRMATION_PERIOD)
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Check and done.")
         self.check_state((self.checkers.crossystem_checker, {
diff --git a/server/site_tests/firmware_InvalidUSB/firmware_InvalidUSB.py b/server/site_tests/firmware_InvalidUSB/firmware_InvalidUSB.py
index e8262b9..4311e3b 100644
--- a/server/site_tests/firmware_InvalidUSB/firmware_InvalidUSB.py
+++ b/server/site_tests/firmware_InvalidUSB/firmware_InvalidUSB.py
@@ -53,7 +53,7 @@
         self.switcher.bypass_rec_mode()
         logging.info('Wait to ensure the USB image is unable to boot...')
         try:
-            self.wait_for_client()
+            self.switcher.wait_for_client()
             raise error.TestFail('Should not boot from the invalid USB image.')
         except ConnectionError:
             logging.info(
@@ -61,7 +61,7 @@
 
         self.restore_usb()
         self.servo.switch_usbkey('dut')
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Expected to boot the restored USB image and reboot.")
         self.check_state((self.checkers.crossystem_checker, {
diff --git a/server/site_tests/firmware_LegacyRecovery/firmware_LegacyRecovery.py b/server/site_tests/firmware_LegacyRecovery/firmware_LegacyRecovery.py
index 478a9d3..b331a64 100644
--- a/server/site_tests/firmware_LegacyRecovery/firmware_LegacyRecovery.py
+++ b/server/site_tests/firmware_LegacyRecovery/firmware_LegacyRecovery.py
@@ -39,7 +39,7 @@
         self.switcher.mode_aware_reboot(wait_for_dut_up=False)
         self.switcher.bypass_rec_mode()
         try:
-            self.wait_for_client()
+            self.switcher.wait_for_client()
         except ConnectionError:
             raise error.TestError('Failed to boot the USB image.')
         self.faft_client.system.run_shell_command(
@@ -54,12 +54,12 @@
         self.switcher.mode_aware_reboot(wait_for_dut_up=False)
         logging.info('Wait to ensure DUT doesnt Boot on USB at Remove screen.')
         try:
-            self.wait_for_client()
+            self.switcher.wait_for_client()
             raise error.TestFail('Unexpected USB boot at Remove Screen.')
         except ConnectionError:
             logging.info('Done, Waited till timeout and no USB boot occured.')
         self.switcher.bypass_rec_mode()
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Expected to boot the restored USB image and reboot.")
         self.check_state((self.checkers.crossystem_checker, {
diff --git a/server/site_tests/firmware_SelfSignedBoot/firmware_SelfSignedBoot.py b/server/site_tests/firmware_SelfSignedBoot/firmware_SelfSignedBoot.py
index c28cf18..446e5c1 100644
--- a/server/site_tests/firmware_SelfSignedBoot/firmware_SelfSignedBoot.py
+++ b/server/site_tests/firmware_SelfSignedBoot/firmware_SelfSignedBoot.py
@@ -3,7 +3,6 @@
 # found in the LICENSE file.
 
 import logging
-import time
 
 from autotest_lib.client.common_lib import error
 from autotest_lib.server.cros import vboot_constants as vboot
@@ -107,7 +106,7 @@
         self.resignimage_ssdkeys()
         self.switcher.mode_aware_reboot(wait_for_dut_up=False)
         self.switcher.bypass_dev_boot_usb()
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Expected USB boot.")
         self.check_state((self.checkers.dev_boot_usb_checker, True,
diff --git a/server/site_tests/firmware_SoftwareSync/firmware_SoftwareSync.py b/server/site_tests/firmware_SoftwareSync/firmware_SoftwareSync.py
index 35da050..801191e 100644
--- a/server/site_tests/firmware_SoftwareSync/firmware_SoftwareSync.py
+++ b/server/site_tests/firmware_SoftwareSync/firmware_SoftwareSync.py
@@ -57,7 +57,7 @@
         self.record_hash_and_corrupt()
         self.sync_and_ec_reboot()
         self.wait_software_sync_and_boot()
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Expect EC in RW and RW is restored.")
         self.check_state(self.software_sync_checker)
diff --git a/server/site_tests/firmware_TPMKernelVersion/firmware_TPMKernelVersion.py b/server/site_tests/firmware_TPMKernelVersion/firmware_TPMKernelVersion.py
index b188cfb..ad4fc66 100644
--- a/server/site_tests/firmware_TPMKernelVersion/firmware_TPMKernelVersion.py
+++ b/server/site_tests/firmware_TPMKernelVersion/firmware_TPMKernelVersion.py
@@ -54,7 +54,7 @@
         self.faft_client.system.set_dev_boot_usb(1)
         self.switcher.mode_aware_reboot(wait_for_dut_up=False)
         self.switcher.bypass_dev_boot_usb()
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         # Check that DUT is booted from USB.
         self.check_state((self.checkers.crossystem_checker,
diff --git a/server/site_tests/firmware_UserRequestRecovery/firmware_UserRequestRecovery.py b/server/site_tests/firmware_UserRequestRecovery/firmware_UserRequestRecovery.py
index 42cce23..400b689 100644
--- a/server/site_tests/firmware_UserRequestRecovery/firmware_UserRequestRecovery.py
+++ b/server/site_tests/firmware_UserRequestRecovery/firmware_UserRequestRecovery.py
@@ -48,7 +48,7 @@
         self.faft_client.system.request_recovery_boot()
         self.switcher.mode_aware_reboot(wait_for_dut_up=False)
         self.switcher.bypass_rec_mode()
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Expected recovery boot, request recovery again.")
         self.check_state((self.checkers.crossystem_checker, {
@@ -59,7 +59,7 @@
         self.switcher.mode_aware_reboot(wait_for_dut_up=False)
         if not dev_mode:
             self.switcher.bypass_rec_mode()
-        self.wait_for_client()
+        self.switcher.wait_for_client()
 
         logging.info("Expected recovery boot.")
         self.check_state((self.checkers.crossystem_checker, {