cr50_test: wait for reboot after updating cr50

After usb_update has run, wait for cr50 to reboot before trying to send
anymore commands.

BUG=none
BRANCH=none
TEST=none

Change-Id: I58db4c77bc02aab0425037a5ba5c2e19d310cc61
Signed-off-by: Mary Ruthven <mruthven@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/567429
Reviewed-by: Aseda Aboagye <aaboagye@chromium.org>
diff --git a/server/cros/faft/cr50_test.py b/server/cros/faft/cr50_test.py
index c0a268c..710d493 100644
--- a/server/cros/faft/cr50_test.py
+++ b/server/cros/faft/cr50_test.py
@@ -182,6 +182,10 @@
 
         rw_ver = self._cr50_run_update(path)
 
+        # Running the update may cause cr50 to reboot. Wait for that before
+        # sending more commands
+        self.cr50.wait_for_reboot()
+
         if erase_nvmem and rollback:
             self.cr50.erase_nvmem()
 
diff --git a/server/cros/servo/chrome_cr50.py b/server/cros/servo/chrome_cr50.py
index 26594c0..6f92711 100644
--- a/server/cros/servo/chrome_cr50.py
+++ b/server/cros/servo/chrome_cr50.py
@@ -40,6 +40,7 @@
     FWMP_LOCKED_PROD = ["Managed device console can't be unlocked"]
     FWMP_LOCKED_DBG = ['Ignoring FWMP unlock setting']
     MAX_RETRY_COUNT = 5
+    START_STR = ['(.*Console is enabled;)']
 
 
     def __init__(self, servo):
@@ -105,12 +106,26 @@
 
     def reboot(self):
         """Reboot Cr50 and wait for CCD to be enabled"""
+        self.send_command('reboot')
+        self.wait_for_reboot()
+
+
+    def wait_for_reboot(self, timeout=60):
+        """Wait for cr50 to reboot"""
         if self.using_ccd():
-            self.send_command('reboot')
-            self.wait_for_ccd_disable()
+            # Cr50 USB is reset when it reboots. Wait for the CCD connection to
+            # go down to detect the reboot.
+            self.wait_for_ccd_disable(timeout, raise_error=False)
             self.ccd_enable()
         else:
-            self.send_command_get_output('reboot', ['Console is enabled;'])
+            # Look for the boot string declaring the console is ready. If we
+            # don't find it, its ok. The command will timeout after 3 seconds
+            # which is longer than the time it takes for cr50 to reboot.
+            try:
+                rv = self.send_command_get_output('\n\n', self.START_STR)
+                logging.debug(rv[0][0])
+            except:
+                pass
 
 
     def rollback(self, eraseflashinfo=True):