faft: Allow enable_rec_mode_and_reboot() to switch USB state

Several helper functions that wrap enable_rec_mode_and_reboot() (most
importantly enable_dev_mode_and_reboot()) change the USB state before
calling it to make the DUT either boot right into recovery mode from
that stick or ensure that it will halt at the recovery screen. This is a
problem when these functions are called while the DUT is already booted
from USB, since enable_rec_mode_and_reboot() tries to sync the root file
system before actually rebooting -- if you have already pulled out the
USB stick from under it at this point, this sync will fail and result in
exceptions.

This change adds an optional usb_state parameter to
enable_rec_mode_and_reboot() and makes it do the switch itself at the
right time (after the DUT is powered off) when provided. All helper
functions calling it in this manner are updated to use the new
functionality.

BUG=None
TEST=Ran firmware_RecoveryButton on a Jerry.

Change-Id: I0420af8b86d6995ad9cd93d2cda5fead1f421687
Signed-off-by: Julius Werner <jwerner@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/251053
Reviewed-by: Yusuf Mohsinally <mohsinally@chromium.org>
diff --git a/server/cros/faft/firmware_test.py b/server/cros/faft/firmware_test.py
index 5d735c2..662b9bf 100644
--- a/server/cros/faft/firmware_test.py
+++ b/server/cros/faft/firmware_test.py
@@ -318,8 +318,7 @@
         @raise TestError: if failed to boot the USB image.
         """
         logging.info('Try boot into USB image...')
-        self.servo.switch_usbkey('host')
-        self.enable_rec_mode_and_reboot()
+        self.enable_rec_mode_and_reboot('host')
         self.wait_fw_screen_and_plug_usb()
         try:
             self.wait_for_client(install_deps=True)
@@ -931,7 +930,7 @@
         """Power cycle DUT AC power."""
         self._client.power_cycle(self.power_control)
 
-    def enable_rec_mode_and_reboot(self):
+    def enable_rec_mode_and_reboot(self, usb_state=None):
         """Switch to rec mode and reboot.
 
         This method emulates the behavior of the old physical recovery switch,
@@ -941,6 +940,8 @@
         self.blocking_sync()
         psc = self.servo.get_power_state_controller()
         psc.power_off()
+        if usb_state:
+            self.servo.switch_usbkey(usb_state)
         psc.power_on(psc.REC_ON)
 
     def enable_dev_mode_and_reboot(self):
@@ -986,10 +987,9 @@
     def enable_keyboard_dev_mode(self):
         """Enable keyboard controlled developer mode"""
         logging.info("Enabling keyboard controlled developer mode")
-        # Plug out USB disk for preventing recovery boot without warning
-        self.servo.switch_usbkey('host')
         # Rebooting EC with rec mode on. Should power on AP.
-        self.enable_rec_mode_and_reboot()
+        # Plug out USB disk for preventing recovery boot without warning
+        self.enable_rec_mode_and_reboot(usb_state='host')
         self.wait_for_client_offline()
         self.wait_fw_screen_and_switch_keyboard_dev_mode(dev=True)
 
@@ -1257,11 +1257,10 @@
         shim to reset TPM values.
         """
         # Unplug USB first to avoid the complicated USB autoboot cases.
-        self.servo.switch_usbkey('host')
         is_dev = self.checkers.crossystem_checker({'devsw_boot': '1'})
         if not is_dev:
             self.enable_dev_mode_and_reboot()
-        self.enable_rec_mode_and_reboot()
+        self.enable_rec_mode_and_reboot(usb_state='host')
         self.wait_fw_screen_and_plug_usb()
         time.sleep(self.faft_config.install_shim_done)
         self.reboot_warm_trigger()