diff --git a/server/cros/faftsequence.py b/server/cros/faftsequence.py
index 49b9485..99b4078 100644
--- a/server/cros/faftsequence.py
+++ b/server/cros/faftsequence.py
@@ -272,10 +272,9 @@
         if not self.root_part_checker(part):
             self.copy_kernel_and_rootfs(from_part=self.OTHER_KERNEL_MAP[part],
                                         to_part=part)
-            self.reset_and_prioritize_kernel(part)
-            self.sync_and_hw_reboot()
-            self.wait_for_client_offline()
-            self.wait_for_client()
+            self.run_faft_step({
+                'userspace_action': (self.reset_and_prioritize_kernel, part),
+            })
 
 
     def wait_fw_screen_and_ctrl_d(self):
@@ -341,7 +340,8 @@
                 logging.info('System is not in dev mode. Reboot into it.')
                 self.run_faft_step({
                     'userspace_action': (self.faft_client.run_shell_command,
-                        'chromeos-firmwareupdate --mode todev && reboot')
+                        'chromeos-firmwareupdate --mode todev && reboot'),
+                    'reboot_action': None,
                 })
         else:
             if not self.crossystem_checker({'devsw_cur': '0'}):
@@ -352,7 +352,8 @@
                 logging.info('System is not in normal mode. Reboot into it.')
                 self.run_faft_step({
                     'userspace_action': (self.faft_client.run_shell_command,
-                        'chromeos-firmwareupdate --mode tonormal && reboot')
+                        'chromeos-firmwareupdate --mode tonormal && reboot'),
+                    'reboot_action': None,
                 })
 
 
diff --git a/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py b/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py
index d85ca1e..c8b9202 100644
--- a/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py
+++ b/server/site_tests/firmware_CorruptKernelA/firmware_CorruptKernelA.py
@@ -2,7 +2,6 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-from autotest_lib.client.common_lib import error
 from autotest_lib.server.cros.faftsequence import FAFTSequence
 
 
diff --git a/server/site_tests/firmware_CorruptKernelB/firmware_CorruptKernelB.py b/server/site_tests/firmware_CorruptKernelB/firmware_CorruptKernelB.py
index 123cb93..38b4e3d 100644
--- a/server/site_tests/firmware_CorruptKernelB/firmware_CorruptKernelB.py
+++ b/server/site_tests/firmware_CorruptKernelB/firmware_CorruptKernelB.py
@@ -2,9 +2,6 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-import time
-
-from autotest_lib.client.common_lib import error
 from autotest_lib.server.cros.faftsequence import FAFTSequence