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']):
diff --git a/server/site_tests/firmware_DevTriggerRecovery/firmware_DevTriggerRecovery.py b/server/site_tests/firmware_DevTriggerRecovery/firmware_DevTriggerRecovery.py
index e6abf4c..26e124a 100644
--- a/server/site_tests/firmware_DevTriggerRecovery/firmware_DevTriggerRecovery.py
+++ b/server/site_tests/firmware_DevTriggerRecovery/firmware_DevTriggerRecovery.py
@@ -22,18 +22,6 @@
     need_dev_transition = False
 
 
-    def wait_fw_screen_and_trigger_recovery(self):
-        """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 self.need_dev_transition:
-            time.sleep(self.TEXT_SCREEN_DELAY)
-            self.send_ctrl_d_to_dut()
-
-
     # The devsw off->on transition states are different based on platforms.
     # For Alex/ZGB, it is dev switch on but normal firmware boot.
     # For other platforms, it is dev switch on and developer firmware boot.
@@ -99,7 +87,8 @@
                 # Ignore the default reboot_action here because the
                 # userspace_action (firmware updater) will reboot the system.
                 'reboot_action': None,
-                'firmware_action': self.wait_fw_screen_and_trigger_recovery,
+                'firmware_action': (self.wait_fw_screen_and_trigger_recovery,
+                                    self.need_dev_transition),
                 'install_deps_after_boot': True,
             },
             {   # Step 3, expected recovery boot and disable dev switch
diff --git a/server/site_tests/firmware_FwScreenPressPower/control b/server/site_tests/firmware_FwScreenPressPower/control
new file mode 100644
index 0000000..b90ddbe
--- /dev/null
+++ b/server/site_tests/firmware_FwScreenPressPower/control
@@ -0,0 +1,26 @@
+# Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+AUTHOR = "Chrome OS Team"
+NAME = "firmware_FwScreenPressPower"
+PURPOSE = "Servo based power button triggered shutdown during firmware screens."
+CRITERIA = "This test will fail if DUT doesn't shutdown"
+TIME = "LONG"
+TEST_CATEGORY = "Functional"
+TEST_CLASS = "firmware"
+TEST_TYPE = "server"
+
+DOC = """
+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.
+"""
+
+def run_fwscreenpresspower(machine):
+    host = hosts.create_host(machine)
+    job.run_test("firmware_FwScreenPressPower", host=host, cmdline_args=args,
+                 use_faft=True, disable_sysinfo=True)
+
+parallel_simple(run_fwscreenpresspower, machines)
diff --git a/server/site_tests/firmware_FwScreenPressPower/firmware_FwScreenPressPower.py b/server/site_tests/firmware_FwScreenPressPower/firmware_FwScreenPressPower.py
new file mode 100644
index 0000000..d4d05d5
--- /dev/null
+++ b/server/site_tests/firmware_FwScreenPressPower/firmware_FwScreenPressPower.py
@@ -0,0 +1,127 @@
+# Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
+# 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.server.cros.faftsequence import FAFTSequence
+
+
+class firmware_FwScreenPressPower(FAFTSequence):
+    """
+    Servo based power button triggered shutdown test 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.
+    """
+    version = 1
+
+
+    def wait_insert_screen_and_press_power(self):
+        """Wait and trigger recovery insert screen and press power button."""
+        self.wait_fw_screen_and_trigger_recovery()
+        self.wait_fw_screen_and_press_power()
+
+
+    def wait_yuck_screen_and_press_power(self):
+        """Wait and trigger yuck screen and press power button."""
+        self.wait_fw_screen_and_trigger_recovery()
+        # Insert a corrupted USB stick. A yuck screen is expected.
+        # This USB stick will be removed in cleanup phase.
+        time.sleep(self.FIRMWARE_SCREEN_DELAY)
+        self.servo.set('usb_mux_sel1', 'dut_sees_usbkey')
+        self.wait_fw_screen_and_press_power()
+
+
+    def setup(self):
+        super(firmware_FwScreenPressPower, self).setup()
+        self.setup_dev_mode(dev_mode=False)
+        self.servo.set('usb_mux_sel1', 'servo_sees_usbkey')
+        usb_dev = self.servo.probe_host_usb_dev()
+        # Corrupt the kernel of USB stick. It is needed for triggering a
+        # yuck screen later.
+        self.corrupt_usb_kernel(usb_dev)
+
+
+    def cleanup(self):
+        self.servo.set('usb_mux_sel1', 'servo_sees_usbkey')
+        usb_dev = self.servo.probe_host_usb_dev()
+        # Restore the kernel of USB stick which is corrupted on setup phase.
+        self.restore_usb_kernel(usb_dev)
+        super(firmware_FwScreenPressPower, self).cleanup()
+
+
+    def run_once(self, host=None):
+        self.register_faft_sequence((
+            {   # Step 1, enable dev mode and dev firmware. When the developer
+                # screen shown, press power button to make DUT shutdown.
+                'state_checker': (self.crossystem_checker, {
+                    'devsw_boot': '0',
+                    'mainfw_type': 'normal',
+                    'recoverysw_boot': '0',
+                }),
+                'userspace_action': self.enable_dev_mode_and_fw,
+                'reboot_action': None,
+                'firmware_action': (self.run_shutdown_process,
+                                    self.wait_fw_screen_and_press_power,
+                                    None,
+                                    self.wait_fw_screen_and_ctrl_d),
+            },
+            {   # Step 2, reboot. When the developer screen shown, press
+                # enter key to trigger recovery insert screen. Then press
+                # power button to make DUT shutdown.
+                'state_checker': (self.crossystem_checker, {
+                    'devsw_boot': '1',
+                    'mainfw_type': 'developer',
+                    'recoverysw_boot': '0',
+                }),
+                'firmware_action': (self.run_shutdown_process,
+                                    self.wait_insert_screen_and_press_power,
+                                    None,
+                                    self.wait_fw_screen_and_ctrl_d),
+            },
+            {   # Step 3, reboot. When the developer screen shown, press
+                # enter key and insert a corrupted USB stick to trigger
+                # yuck screen. Then press power button to make DUT shutdown.
+                'state_checker': (self.crossystem_checker, {
+                    'devsw_boot': '1',
+                    'mainfw_type': 'developer',
+                    'recoverysw_boot': '0',
+                }),
+                'firmware_action': (self.run_shutdown_process,
+                                    self.wait_yuck_screen_and_press_power,
+                                    None,
+                                    self.wait_fw_screen_and_ctrl_d),
+            },
+            {   # Step 4, enable normal mode and normal firmware.
+                'state_checker': (self.crossystem_checker, {
+                    'devsw_boot': '1',
+                    'mainfw_type': 'developer',
+                    'recoverysw_boot': '0',
+                }),
+                'userspace_action': self.enable_normal_mode_and_fw,
+                'reboot_action': None,
+            },
+            {   # Step 5, turn on the recovery boot. Since a USB stick was
+                # inserted in step 3, a recovery remove screen is shown.
+                # Press power button to make DUT shutdown.
+                'state_checker': (self.crossystem_checker, {
+                    'devsw_boot': '0',
+                    'mainfw_type': 'normal',
+                    'recoverysw_boot': '0',
+                }),
+                'userspace_action': self.faft_client.request_recovery_boot,
+                'firmware_action': (self.run_shutdown_process,
+                                    self.wait_fw_screen_and_press_power),
+            },
+            {   # Step 6, done.
+                'state_checker': (self.crossystem_checker, {
+                    'devsw_boot': '0',
+                    'mainfw_type': 'normal',
+                    'recoverysw_boot': '0',
+                }),
+            },
+        ))
+        self.run_faft_sequence()
diff --git a/server/site_tests/suites/control.faft b/server/site_tests/suites/control.faft
index 408f278..a20d6e0 100644
--- a/server/site_tests/suites/control.faft
+++ b/server/site_tests/suites/control.faft
@@ -18,6 +18,7 @@
 TESTS_ON_NORMAL_MODE = [
     ('firmware_DevMode', {}),
     ('firmware_DevTriggerRecovery', {}),
+    ('firmware_FwScreenPressPower', {}),
 ]
 
 TESTS_ON_DEV_MODE = [