faft: Unify all USB disk setup code in a new method setup_usbkey()

This CL adds a new method setup_usbkey() in FAFTSequence class, which
helps all USB disk related setup and check. Its usage:

If the test requires USB key (Chrome OS test image inside) and muxes to host:
    self.setup_usbkey(usbkey=True, host=True)

If the test requires USB key (Chrome OS test image inside) and muxes to DUT:
    self.setup_usbkey(usbkey=True, host=False)

If the test does not requires USB key:
    self.setup_usbkey(usbkey=False)

BUG=chromium-os:35902
TEST=run the FAFT suite:
$ run_remote_tests.sh --board link --remote dut suite:faft_bios

Change-Id: If6eb94ee52f612df568c4f2f07a7acbbdd0a0b66
Reviewed-on: https://gerrit.chromium.org/gerrit/37193
Reviewed-by: Vic Yang <victoryang@chromium.org>
Commit-Ready: Tom Wai-Hong Tam <waihong@chromium.org>
Tested-by: Tom Wai-Hong Tam <waihong@chromium.org>
diff --git a/server/cros/faftsequence.py b/server/cros/faftsequence.py
index cde9259..3cea48e 100644
--- a/server/cros/faftsequence.py
+++ b/server/cros/faftsequence.py
@@ -398,6 +398,30 @@
         self.mark_setup_done('usb_check')
 
 
+    def setup_usbkey(self, usbkey, host=None):
+        """Setup the USB disk for the test.
+
+        It checks the setup of USB disk and a valid ChromeOS test image inside.
+        It also muxes the USB disk to either the host or DUT by request.
+
+        Args:
+          usbkey: True if the USB disk is required for the test, False if not
+                  required.
+          host: Optional, True to mux the USB disk to host, False to mux it
+                to DUT, default to do nothing.
+        """
+        if usbkey:
+            self.assert_test_image_in_usb_disk()
+        elif host is None:
+            # USB disk is not required for the test. Better to mux it to host.
+            host = True
+
+        if host is True:
+            self.servo.set('usb_mux_sel1', 'servo_sees_usbkey')
+        elif host is False:
+            self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+
+
     def get_server_address(self):
         """Get the server address seen from the client.
 
diff --git a/server/site_tests/firmware_CgptState/firmware_CgptState.py b/server/site_tests/firmware_CgptState/firmware_CgptState.py
index 860cab4..cdd2305 100644
--- a/server/site_tests/firmware_CgptState/firmware_CgptState.py
+++ b/server/site_tests/firmware_CgptState/firmware_CgptState.py
@@ -45,6 +45,7 @@
         super(firmware_CgptState, self).setup()
         self.host = host
         self.setup_dev_mode(dev_mode=False)
+        self.setup_usbkey(usbkey=False)
         self.setup_kernel('a')
 
 
diff --git a/server/site_tests/firmware_CgptStress/firmware_CgptStress.py b/server/site_tests/firmware_CgptStress/firmware_CgptStress.py
index 324ebd3..9083abc 100644
--- a/server/site_tests/firmware_CgptStress/firmware_CgptStress.py
+++ b/server/site_tests/firmware_CgptStress/firmware_CgptStress.py
@@ -27,6 +27,7 @@
     def setup(self, dev_mode=False):
         super(firmware_CgptStress, self).setup()
         self.setup_dev_mode(dev_mode)
+        self.setup_usbkey(usbkey=False)
         self.setup_kernel('a')
 
 
diff --git a/server/site_tests/firmware_ConsecutiveBoot/firmware_ConsecutiveBoot.py b/server/site_tests/firmware_ConsecutiveBoot/firmware_ConsecutiveBoot.py
index 85f7f36..5621f0b 100644
--- a/server/site_tests/firmware_ConsecutiveBoot/firmware_ConsecutiveBoot.py
+++ b/server/site_tests/firmware_ConsecutiveBoot/firmware_ConsecutiveBoot.py
@@ -32,6 +32,7 @@
     def setup(self, dev_mode=False):
         super(firmware_ConsecutiveBoot, self).setup()
         self.setup_dev_mode(dev_mode)
+        self.setup_usbkey(usbkey=False)
 
 
     def run_once(self, dev_mode=False, host=None):
diff --git a/server/site_tests/firmware_CorruptBothFwBodyAB/firmware_CorruptBothFwBodyAB.py b/server/site_tests/firmware_CorruptBothFwBodyAB/firmware_CorruptBothFwBodyAB.py
index 4dd3ee1..4909a2b 100644
--- a/server/site_tests/firmware_CorruptBothFwBodyAB/firmware_CorruptBothFwBodyAB.py
+++ b/server/site_tests/firmware_CorruptBothFwBodyAB/firmware_CorruptBothFwBodyAB.py
@@ -33,9 +33,8 @@
             self.use_ro = True
             self.setup_dev_mode(dev_mode)
         else:
-            self.assert_test_image_in_usb_disk()
             self.setup_dev_mode(dev_mode)
-            self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+            self.setup_usbkey(usbkey=True, host=False)
 
 
     def cleanup(self):
diff --git a/server/site_tests/firmware_CorruptBothFwSigAB/firmware_CorruptBothFwSigAB.py b/server/site_tests/firmware_CorruptBothFwSigAB/firmware_CorruptBothFwSigAB.py
index 66534c3..31c0948 100644
--- a/server/site_tests/firmware_CorruptBothFwSigAB/firmware_CorruptBothFwSigAB.py
+++ b/server/site_tests/firmware_CorruptBothFwSigAB/firmware_CorruptBothFwSigAB.py
@@ -22,9 +22,8 @@
     def setup(self, dev_mode=False):
         super(firmware_CorruptBothFwSigAB, self).setup()
         self.backup_firmware()
-        self.assert_test_image_in_usb_disk()
         self.setup_dev_mode(dev_mode)
-        self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+        self.setup_usbkey(usbkey=True, host=False)
 
 
     def cleanup(self):
diff --git a/server/site_tests/firmware_CorruptBothKernelAB/firmware_CorruptBothKernelAB.py b/server/site_tests/firmware_CorruptBothKernelAB/firmware_CorruptBothKernelAB.py
index cf95daf..b90c831 100644
--- a/server/site_tests/firmware_CorruptBothKernelAB/firmware_CorruptBothKernelAB.py
+++ b/server/site_tests/firmware_CorruptBothKernelAB/firmware_CorruptBothKernelAB.py
@@ -38,10 +38,9 @@
 
     def setup(self, dev_mode=False):
         super(firmware_CorruptBothKernelAB, self).setup()
-        self.assert_test_image_in_usb_disk()
         self.setup_dev_mode(dev_mode)
+        self.setup_usbkey(usbkey=True, host=False)
         self.setup_kernel('a')
-        self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
 
 
     def cleanup(self):
diff --git a/server/site_tests/firmware_CorruptFwBodyA/firmware_CorruptFwBodyA.py b/server/site_tests/firmware_CorruptFwBodyA/firmware_CorruptFwBodyA.py
index 861146d..1b0a6f1 100644
--- a/server/site_tests/firmware_CorruptFwBodyA/firmware_CorruptFwBodyA.py
+++ b/server/site_tests/firmware_CorruptFwBodyA/firmware_CorruptFwBodyA.py
@@ -25,6 +25,7 @@
         super(firmware_CorruptFwBodyA, self).setup()
         self.backup_firmware()
         self.setup_dev_mode(dev_mode)
+        self.setup_usbkey(usbkey=False)
 
 
     def cleanup(self):
diff --git a/server/site_tests/firmware_CorruptFwBodyB/firmware_CorruptFwBodyB.py b/server/site_tests/firmware_CorruptFwBodyB/firmware_CorruptFwBodyB.py
index 5b18dcf..e1e4a6a 100644
--- a/server/site_tests/firmware_CorruptFwBodyB/firmware_CorruptFwBodyB.py
+++ b/server/site_tests/firmware_CorruptFwBodyB/firmware_CorruptFwBodyB.py
@@ -23,6 +23,7 @@
         super(firmware_CorruptFwBodyB, self).setup()
         self.backup_firmware()
         self.setup_dev_mode(dev_mode)
+        self.setup_usbkey(usbkey=False)
 
 
     def cleanup(self):
diff --git a/server/site_tests/firmware_CorruptFwSigA/firmware_CorruptFwSigA.py b/server/site_tests/firmware_CorruptFwSigA/firmware_CorruptFwSigA.py
index 488ba89..d39ca7d 100644
--- a/server/site_tests/firmware_CorruptFwSigA/firmware_CorruptFwSigA.py
+++ b/server/site_tests/firmware_CorruptFwSigA/firmware_CorruptFwSigA.py
@@ -16,6 +16,7 @@
         super(firmware_CorruptFwSigA, self).setup()
         self.backup_firmware()
         self.setup_dev_mode(dev_mode)
+        self.setup_usbkey(usbkey=False)
 
 
     def cleanup(self):
diff --git a/server/site_tests/firmware_CorruptFwSigB/firmware_CorruptFwSigB.py b/server/site_tests/firmware_CorruptFwSigB/firmware_CorruptFwSigB.py
index e038d84..6c91246 100644
--- a/server/site_tests/firmware_CorruptFwSigB/firmware_CorruptFwSigB.py
+++ b/server/site_tests/firmware_CorruptFwSigB/firmware_CorruptFwSigB.py
@@ -16,6 +16,7 @@
         super(firmware_CorruptFwSigB, self).setup()
         self.backup_firmware()
         self.setup_dev_mode(dev_mode)
+        self.setup_usbkey(usbkey=False)
 
 
     def cleanup(self):
diff --git a/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py b/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py
index 579f5ce..69bbeec 100644
--- a/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py
+++ b/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py
@@ -18,6 +18,7 @@
     def setup(self, dev_mode=False):
         super(firmware_CorruptKernelA, self).setup()
         self.setup_dev_mode(dev_mode)
+        self.setup_usbkey(usbkey=False)
         self.setup_kernel('a')
 
 
diff --git a/server/site_tests/firmware_CorruptKernelB/firmware_CorruptKernelB.py b/server/site_tests/firmware_CorruptKernelB/firmware_CorruptKernelB.py
index 2cb6ce3..e0d4059 100644
--- a/server/site_tests/firmware_CorruptKernelB/firmware_CorruptKernelB.py
+++ b/server/site_tests/firmware_CorruptKernelB/firmware_CorruptKernelB.py
@@ -19,6 +19,7 @@
     def setup(self, dev_mode=False):
         super(firmware_CorruptKernelB, self).setup()
         self.setup_dev_mode(dev_mode)
+        self.setup_usbkey(usbkey=False)
         self.setup_kernel('a')
 
 
diff --git a/server/site_tests/firmware_DevBootUSB/firmware_DevBootUSB.py b/server/site_tests/firmware_DevBootUSB/firmware_DevBootUSB.py
index 1bb37b1..f2d8c3a 100644
--- a/server/site_tests/firmware_DevBootUSB/firmware_DevBootUSB.py
+++ b/server/site_tests/firmware_DevBootUSB/firmware_DevBootUSB.py
@@ -23,9 +23,8 @@
 
     def setup(self):
         super(firmware_DevBootUSB, self).setup()
-        self.assert_test_image_in_usb_disk()
         self.setup_dev_mode(dev_mode=True)
-        self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+        self.setup_usbkey(usbkey=True, host=False)
 
         self.original_dev_boot_usb = self.faft_client.get_dev_boot_usb()
         logging.info('Original dev_boot_usb value: %s',
diff --git a/server/site_tests/firmware_DevFwNormalBoot/firmware_DevFwNormalBoot.py b/server/site_tests/firmware_DevFwNormalBoot/firmware_DevFwNormalBoot.py
index 387b5a5..4d4faf8 100644
--- a/server/site_tests/firmware_DevFwNormalBoot/firmware_DevFwNormalBoot.py
+++ b/server/site_tests/firmware_DevFwNormalBoot/firmware_DevFwNormalBoot.py
@@ -40,10 +40,9 @@
         # This test is only meaningful on Alex/ZGB.
         if self.faft_client.get_platform_name() in ('Alex', 'ZGB'):
             self.has_different_dev_fw = True
-            self.assert_test_image_in_usb_disk()
-            self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
             # This test is run on developer mode only.
             self.setup_dev_mode(dev_mode=True)
+            self.setup_usbkey(usbkey=True, host=False)
 
 
     def run_once(self, host=None):
diff --git a/server/site_tests/firmware_DevMode/firmware_DevMode.py b/server/site_tests/firmware_DevMode/firmware_DevMode.py
index 6cd67fa..b8c6a07 100644
--- a/server/site_tests/firmware_DevMode/firmware_DevMode.py
+++ b/server/site_tests/firmware_DevMode/firmware_DevMode.py
@@ -15,6 +15,7 @@
     def setup(self):
         super(firmware_DevMode, self).setup()
         self.setup_dev_mode(dev_mode=False)
+        self.setup_usbkey(usbkey=False)
 
 
     def run_once(self, host=None):
diff --git a/server/site_tests/firmware_DevModeStress/firmware_DevModeStress.py b/server/site_tests/firmware_DevModeStress/firmware_DevModeStress.py
index a83026c..cee3b71 100644
--- a/server/site_tests/firmware_DevModeStress/firmware_DevModeStress.py
+++ b/server/site_tests/firmware_DevModeStress/firmware_DevModeStress.py
@@ -59,6 +59,11 @@
         super(firmware_DevModeStress, self).initialize(host, cmdline_args,
                                                        use_pyauto, use_faft)
 
+    def setup(self):
+        super(firmware_DevModeStress, self).setup()
+        self.setup_usbkey(usbkey=False)
+
+
     def run_once(self, host=None):
         self.register_faft_sequence((
             {   # Step 1, verify dev mode
diff --git a/server/site_tests/firmware_DevScreenTimeout/firmware_DevScreenTimeout.py b/server/site_tests/firmware_DevScreenTimeout/firmware_DevScreenTimeout.py
index 618e7d7..3548b70 100644
--- a/server/site_tests/firmware_DevScreenTimeout/firmware_DevScreenTimeout.py
+++ b/server/site_tests/firmware_DevScreenTimeout/firmware_DevScreenTimeout.py
@@ -86,6 +86,7 @@
         super(firmware_DevScreenTimeout, self).setup()
         # This test is run on developer mode only.
         self.setup_dev_mode(dev_mode=True)
+        self.setup_usbkey(usbkey=False)
 
 
     def run_once(self, host=None):
diff --git a/server/site_tests/firmware_DevTriggerRecovery/firmware_DevTriggerRecovery.py b/server/site_tests/firmware_DevTriggerRecovery/firmware_DevTriggerRecovery.py
index 837979a..79898d9 100644
--- a/server/site_tests/firmware_DevTriggerRecovery/firmware_DevTriggerRecovery.py
+++ b/server/site_tests/firmware_DevTriggerRecovery/firmware_DevTriggerRecovery.py
@@ -62,9 +62,8 @@
 
     def setup(self):
         super(firmware_DevTriggerRecovery, self).setup()
-        self.assert_test_image_in_usb_disk()
         self.setup_dev_mode(dev_mode=False)
-        self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+        self.setup_usbkey(usbkey=True, host=False)
 
 
     def run_once(self, host=None):
diff --git a/server/site_tests/firmware_RONormalBoot/firmware_RONormalBoot.py b/server/site_tests/firmware_RONormalBoot/firmware_RONormalBoot.py
index 9a6ff33..a758aec 100644
--- a/server/site_tests/firmware_RONormalBoot/firmware_RONormalBoot.py
+++ b/server/site_tests/firmware_RONormalBoot/firmware_RONormalBoot.py
@@ -27,6 +27,7 @@
         super(firmware_RONormalBoot, self).setup()
         self.backup_firmware()
         self.setup_dev_mode(dev_mode)
+        self.setup_usbkey(usbkey=False)
 
 
     def cleanup(self):
diff --git a/server/site_tests/firmware_RecoveryButton/firmware_RecoveryButton.py b/server/site_tests/firmware_RecoveryButton/firmware_RecoveryButton.py
index 46d8985..1e011c3 100644
--- a/server/site_tests/firmware_RecoveryButton/firmware_RecoveryButton.py
+++ b/server/site_tests/firmware_RecoveryButton/firmware_RecoveryButton.py
@@ -33,9 +33,8 @@
 
     def setup(self, dev_mode=False):
         super(firmware_RecoveryButton, self).setup()
-        self.assert_test_image_in_usb_disk()
         self.setup_dev_mode(dev_mode)
-        self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+        self.setup_usbkey(usbkey=True, host=False)
 
 
     def cleanup(self):
diff --git a/server/site_tests/firmware_RollbackFirmware/firmware_RollbackFirmware.py b/server/site_tests/firmware_RollbackFirmware/firmware_RollbackFirmware.py
index 8f5857b..d374f03 100644
--- a/server/site_tests/firmware_RollbackFirmware/firmware_RollbackFirmware.py
+++ b/server/site_tests/firmware_RollbackFirmware/firmware_RollbackFirmware.py
@@ -21,9 +21,8 @@
     def setup(self, dev_mode=False):
         super(firmware_RollbackFirmware, self).setup()
         self.backup_firmware()
-        self.assert_test_image_in_usb_disk()
         self.setup_dev_mode(dev_mode)
-        self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+        self.setup_usbkey(usbkey=True, host=False)
 
 
     def cleanup(self):
diff --git a/server/site_tests/firmware_RollbackKernel/firmware_RollbackKernel.py b/server/site_tests/firmware_RollbackKernel/firmware_RollbackKernel.py
index d15b92b..f0cb87c 100644
--- a/server/site_tests/firmware_RollbackKernel/firmware_RollbackKernel.py
+++ b/server/site_tests/firmware_RollbackKernel/firmware_RollbackKernel.py
@@ -38,10 +38,9 @@
 
     def setup(self, dev_mode=False):
         super(firmware_RollbackKernel, self).setup()
-        self.assert_test_image_in_usb_disk()
         self.setup_dev_mode(dev_mode)
+        self.setup_usbkey(usbkey=True, host=False)
         self.setup_kernel('a')
-        self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
 
 
     def cleanup(self):
diff --git a/server/site_tests/firmware_SoftwareSync/firmware_SoftwareSync.py b/server/site_tests/firmware_SoftwareSync/firmware_SoftwareSync.py
index b86904c..93554f4 100644
--- a/server/site_tests/firmware_SoftwareSync/firmware_SoftwareSync.py
+++ b/server/site_tests/firmware_SoftwareSync/firmware_SoftwareSync.py
@@ -30,6 +30,7 @@
         super(firmware_SoftwareSync, self).setup()
         self.backup_firmware()
         self.setup_dev_mode(dev_mode)
+        self.setup_usbkey(usbkey=False)
         self.ensure_rw()
 
 
diff --git a/server/site_tests/firmware_TryFwB/firmware_TryFwB.py b/server/site_tests/firmware_TryFwB/firmware_TryFwB.py
index 04d7dc5..9f8ddfc 100644
--- a/server/site_tests/firmware_TryFwB/firmware_TryFwB.py
+++ b/server/site_tests/firmware_TryFwB/firmware_TryFwB.py
@@ -15,6 +15,7 @@
     def setup(self, dev_mode=False):
         super(firmware_TryFwB, self).setup()
         self.setup_dev_mode(dev_mode)
+        self.setup_usbkey(usbkey=False)
         self.setup_tried_fwb(tried_fwb=False)
 
 
diff --git a/server/site_tests/firmware_UpdateECBin/firmware_UpdateECBin.py b/server/site_tests/firmware_UpdateECBin/firmware_UpdateECBin.py
index be1a130..3d3d7fb 100644
--- a/server/site_tests/firmware_UpdateECBin/firmware_UpdateECBin.py
+++ b/server/site_tests/firmware_UpdateECBin/firmware_UpdateECBin.py
@@ -52,6 +52,7 @@
         super(firmware_UpdateECBin, self).setup()
         self.backup_firmware()
         self.setup_dev_mode(dev_mode)
+        self.setup_usbkey(usbkey=False)
 
         temp_path = self.faft_client.get_temp_path()
         self.faft_client.setup_firmwareupdate_temp_dir()
diff --git a/server/site_tests/firmware_UserRequestRecovery/firmware_UserRequestRecovery.py b/server/site_tests/firmware_UserRequestRecovery/firmware_UserRequestRecovery.py
index e4b5881..5e1651c 100644
--- a/server/site_tests/firmware_UserRequestRecovery/firmware_UserRequestRecovery.py
+++ b/server/site_tests/firmware_UserRequestRecovery/firmware_UserRequestRecovery.py
@@ -32,9 +32,8 @@
 
     def setup(self, dev_mode=False):
         super(firmware_UserRequestRecovery, self).setup()
-        self.assert_test_image_in_usb_disk()
         self.setup_dev_mode(dev_mode)
-        self.servo.enable_usb_hub(host=True)
+        self.setup_usbkey(usbkey=True, host=True)
 
 
     def cleanup(self):