A FAFT test of pressing power button during firmware screens.

This test requires a USB disk plugged-in, which contains a Chrome OS test
image (built by "build_image --test"). On runtime, this test triggers
four firmware screens (developer, remove, insert, and yuck screens), and
then presses the power button in order to power the machine down.

BUG=chromium-os:19710
TEST=run_remote_tests.sh --remote=$REMOTE_IP -a "xml_config=$OVERLAY_XML \
         servo_vid=0x18d1 servo_pid=0x5001" FwScreenShutdown
     run_remote_tests.sh --remote=$REMOTE_IP -a "xml_config=$OVERLAY_XML \
         servo_vid=0x18d1 servo_pid=0x5001" DevTriggerRecovery

Change-Id: Iecb614a05835eca86e4a8407d2bea2270b4d97ac
Reviewed-on: https://gerrit.chromium.org/gerrit/13131
Reviewed-by: Randall Spangler <rspangler@chromium.org>
Reviewed-by: Todd Broch <tbroch@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 fd1e0f4..38270a3 100644
--- a/server/cros/faftsequence.py
+++ b/server/cros/faftsequence.py
@@ -76,6 +76,9 @@
     USB_PLUG_DELAY = 10
     SYNC_DELAY = 5
 
+    CHROMEOS_MAGIC = "CHROMEOS"
+    CORRUPTED_MAGIC = "CORRUPTD"
+
     # Recovery reason codes, copied from:
     #     vboot_reference/firmware/lib/vboot_nvstorage.h
     #     vboot_reference/firmware/lib/vboot_struct.h
@@ -409,6 +412,18 @@
         self.send_ctrl_d_to_dut()
 
 
+    def wait_fw_screen_and_trigger_recovery(self, need_dev_transition=False):
+        """Wait for firmware warning screen and trigger recovery boot."""
+        time.sleep(self.FIRMWARE_SCREEN_DELAY)
+        self.send_enter_to_dut()
+
+        # For Alex/ZGB, there is a dev warning screen in text mode.
+        # Skip it by pressing Ctrl-D.
+        if need_dev_transition:
+            time.sleep(self.TEXT_SCREEN_DELAY)
+            self.send_ctrl_d_to_dut()
+
+
     def wait_fw_screen_and_plug_usb(self):
         """Wait for firmware warning screen and then unplug and plug the USB."""
         time.sleep(self.FIRMWARE_SCREEN_DELAY)
@@ -417,6 +432,18 @@
         self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
 
 
+    def wait_fw_screen_and_press_power(self):
+        """Wait for firmware warning screen and press power button."""
+        time.sleep(self.FIRMWARE_SCREEN_DELAY)
+        self.servo.power_normal_press()
+
+
+    def wait_fw_screen_and_close_lid(self):
+        """Wait for firmware warning screen and close lid."""
+        time.sleep(self.FIRMWARE_SCREEN_DELAY)
+        self.servo.lid_close()
+
+
     def setup_tried_fwb(self, tried_fwb):
         """Setup for fw B tried state.
 
@@ -440,6 +467,20 @@
                 self.run_faft_step({})
 
 
+    def enable_dev_mode_and_fw(self):
+        """Enable developer mode and use developer firmware."""
+        self.servo.enable_development_mode()
+        self.faft_client.run_shell_command(
+                'chromeos-firmwareupdate --mode todev && reboot')
+
+
+    def enable_normal_mode_and_fw(self):
+        """Enable normal mode and use normal firmware."""
+        self.servo.disable_development_mode()
+        self.faft_client.run_shell_command(
+                'chromeos-firmwareupdate --mode tonormal && reboot')
+
+
     def setup_dev_mode(self, dev_mode):
         """Setup for development mode.
 
@@ -527,6 +568,61 @@
         self.servo.warm_reset()
 
 
+    def _modify_usb_kernel(self, usb_dev, from_magic, to_magic):
+        """Modify the kernel header magic in USB stick.
+
+        The kernel header magic is the first 8-byte of kernel partition.
+        We modify it to make it fail on kernel verification check.
+
+        Args:
+          usb_dev: A string of USB stick path on the host, like '/dev/sdc'.
+          from_magic: A string of magic which we change it from.
+          to_magic: A string of magic which we change it to.
+
+        Raises:
+          error.TestError: if failed to change magic.
+        """
+        assert len(from_magic) == 8
+        assert len(to_magic) == 8
+        kernel_part = self._join_part(usb_dev, '2')
+        read_cmd = "sudo dd if=%s bs=8 count=1 2>/dev/null" % kernel_part
+        current_magic = utils.system_output(read_cmd)
+        if current_magic == to_magic:
+            logging.info("The kernel magic is already %s." % current_magic)
+            return
+        if current_magic != from_magic:
+            raise error.TestError("Invalid kernel image on USB: wrong magic.")
+
+        logging.info('Modify the kernel magic in USB, from %s to %s.' %
+                     (from_magic, to_magic))
+        write_cmd = ("echo -n '%s' | sudo dd of=%s oflag=sync conv=notrunc "
+                     " 2>/dev/null" % (to_magic, kernel_part))
+        utils.system(write_cmd)
+
+        if utils.system_output(read_cmd) != to_magic:
+            raise error.TestError("Failed to write new magic.")
+
+
+    def corrupt_usb_kernel(self, usb_dev):
+        """Corrupt USB kernel by modifying its magic from CHROMEOS to CORRUPTD.
+
+        Args:
+          usb_dev: A string of USB stick path on the host, like '/dev/sdc'.
+        """
+        self._modify_usb_kernel(usb_dev, self.CHROMEOS_MAGIC,
+                                self.CORRUPTED_MAGIC)
+
+
+    def restore_usb_kernel(self, usb_dev):
+        """Restore USB kernel by modifying its magic from CORRUPTD to CHROMEOS.
+
+        Args:
+          usb_dev: A string of USB stick path on the host, like '/dev/sdc'.
+        """
+        self._modify_usb_kernel(usb_dev, self.CORRUPTED_MAGIC,
+                                self.CHROMEOS_MAGIC)
+
+
     def _call_action(self, action_tuple):
         """Call the action function with/without arguments.
 
@@ -558,6 +654,37 @@
         return None
 
 
+    def run_shutdown_process(self, shutdown_action, pre_power_action=None,
+            post_power_action=None):
+        """Run shutdown_action(), which makes DUT shutdown, and power it on.
+
+        Args:
+          shutdown_action: a function which makes DUT shutdown, like pressing
+                           power key.
+          pre_power_action: a function which is called before next power on.
+          post_power_action: a function which is called after next power on.
+
+        Raises:
+          error.TestFail: if the shutdown_action() failed to turn DUT off.
+        """
+        self._call_action(shutdown_action)
+        logging.info('Wait to ensure DUT shut down...')
+        try:
+            self.wait_for_client()
+            raise error.TestFail(
+                    'Should shut the device down after calling %s.' %
+                    str(shutdown_action))
+        except AssertionError:
+            logging.info(
+                'DUT is surely shutdown. We are going to power it on again...')
+
+        if pre_power_action:
+            self._call_action(pre_power_action)
+        self.servo.power_normal_press()
+        if post_power_action:
+            self._call_action(post_power_action)
+
+
     def register_faft_template(self, template):
         """Register FAFT template, the default FAFT_STEP of each step.
 
@@ -601,7 +728,7 @@
 
         for key in test:
             if key not in FAFT_STEP_KEYS:
-                error.TestError('Invalid key in FAFT step: %s', key)
+                raise error.TestError('Invalid key in FAFT step: %s', key)
 
         if test['state_checker']:
             if not self._call_action(test['state_checker']):