firmware_Cr50WPG3: use ccd to check write protect am: a7b28c535a am: 6b7034ac0a am: 9baa8eee30

Change-Id: I8fec1d41c0ff5e5bf55f255480aa8ee2a866bd4f
diff --git a/server/site_tests/firmware_Cr50WPG3/firmware_Cr50WPG3.py b/server/site_tests/firmware_Cr50WPG3/firmware_Cr50WPG3.py
index 38ae985..e331eb2 100644
--- a/server/site_tests/firmware_Cr50WPG3/firmware_Cr50WPG3.py
+++ b/server/site_tests/firmware_Cr50WPG3/firmware_Cr50WPG3.py
@@ -3,6 +3,7 @@
 # found in the LICENSE file.
 
 import logging
+import re
 import time
 
 from autotest_lib.client.common_lib import error
@@ -14,9 +15,15 @@
     version = 1
 
     WAIT_FOR_STATE = 10
+    WP_REGEX = r'write protect is ((en|dis)abled)\.'
+    FLASHROM_WP_CMD = ('sudo flashrom -p raiden_debug_spi:target=AP,serial=%s '
+                       '--wp-status')
 
     def cleanup(self):
         """Reenable servo wp."""
+        self.cr50.send_command('rddkeepalive disable')
+        self.cr50.send_command('ccdblock IGNORE_SERVO disable')
+        self.servo.enable_main_servo_device()
         try:
             if hasattr(self, '_start_fw_wp_vref'):
                 self.servo.set_nocheck('fw_wp_state', self._start_fw_wp_state)
@@ -25,29 +32,59 @@
             super(firmware_Cr50WPG3, self).cleanup()
 
 
+    def generate_flashrom_wp_cmd(self):
+        """Use the cr50 serialname to generate the flashrom command."""
+        devid = self.servo.get('cr50_devid')
+        serial = devid.replace('0x', '').replace(' ', '-').upper()
+        logging.info('CCD serial: %s', serial)
+        self._flashrom_wp_cmd = self.FLASHROM_WP_CMD % serial
+
+
+    def get_wp_state(self):
+        """Returns 'on' if write protect is enabled. 'off' if it's disabled."""
+        flashrom_output = self.servo.system_output(self._flashrom_wp_cmd)
+        match = re.search(self.WP_REGEX, flashrom_output)
+        logging.info('WP is %s', match.group(1) if match else 'UKNOWN')
+        logging.debug('flashrom output\n%s', flashrom_output)
+        if not match:
+            raise error.TestError('Unable to find WP status in flashrom output')
+        return match.group(1)
+
+
     def run_once(self):
         """Verify WP in G3."""
         if self.check_cr50_capability(['wp_on_in_g3'], suppress_warning=True):
             raise error.TestNAError('WP not pulled up in G3')
-        if self.servo.main_device_is_ccd():
-            raise error.TestNAError('Need servo flex to check WP')
+        if not self.servo.dts_mode_is_valid():
+            raise error.TestNAError('Type-C servo v4 required to check wp.')
+        try:
+            self.cr50.ccd_enable(True)
+        except:
+            raise error.TestNAError('CCD required to check wp.')
+        self.generate_flashrom_wp_cmd()
 
         self.fast_open(True)
+        # faft-cr50 runs with servo micro and type c servo v4. Use ccdblock to
+        # get cr50 to ignore the fact servo is connected and allow the test to
+        # use ccd to check wp status.
+        self.cr50.send_command('ccdblock IGNORE_SERVO enable')
+        self.cr50.send_command('rddkeepalive enable')
+        self.cr50.get_ccdstate()
 
-        self._start_fw_wp_state = self.servo.get('fw_wp_state')
-        self._start_fw_wp_vref = self.servo.get('fw_wp_vref')
-        # Stop forcing wp using servo, so we can set it with ccd.
-        self.servo.set_nocheck('fw_wp_state', 'reset')
-        self.servo.set_nocheck('fw_wp_vref', 'off')
+        if self.servo.main_device_is_flex():
+            self._start_fw_wp_state = self.servo.get('fw_wp_state')
+            self._start_fw_wp_vref = self.servo.get('fw_wp_vref')
+            # Stop forcing wp using servo, so we can set it with ccd.
+            self.servo.set_nocheck('fw_wp_state', 'reset')
+            self.servo.set_nocheck('fw_wp_vref', 'off')
 
         # disable write protect
         self.cr50.send_command('wp disable atboot')
 
         # Verify we can see it's disabled. This should always work. If it
         # doesn't, it may be a setup issue.
-        servo_wp_s0 = self.servo.get('fw_wp_state')
-        logging.info('servo wp: %s', servo_wp_s0)
-        if servo_wp_s0 != 'off':
+        servo_wp_s0 = self.get_wp_state()
+        if servo_wp_s0 == 'enabled':
             raise error.TestError("WP isn't disabled in S0")
 
         self.faft_client.system.run_shell_command('poweroff')
@@ -56,10 +93,9 @@
             self.ec.send_command('hibernate')
             time.sleep(self.WAIT_FOR_STATE)
 
-        servo_wp_g3 = self.servo.get('fw_wp_state')
-        logging.info('servo wp: %s', servo_wp_g3)
+        servo_wp_g3 = self.get_wp_state()
         self._try_to_bring_dut_up()
         # Some boards don't power the EC_WP signal in G3, so it can't be
         # disabled by cr50 in G3.
-        if servo_wp_g3 != 'off':
+        if servo_wp_g3 == 'enabled':
             raise error.TestFail("WP can't be disabled by Cr50 in G3")