[PATCH v5 2/4] test/py: Handle expected reset by command

[Masami Hiramatsu]

Add wait_for_reboot optional argument to ConsoleBase::run_command()
so that it can handle an expected reset by command execution.

This is useful if a command will reset the sandbox while testing
such commands, e.g. run_command("reset", wait_for_reboot = True)

Signed-off-by: Masami Hiramatsu <masami.hiramatsu at linaro.org>
---
 Changes in v5:
  - Split the wait-for-boot-prompt code as an internal function
    so that it is able to be shared with ensure_spawned().
  - Fix wait_for_reboot parameter description.
---
 test/py/u_boot_console_base.py |   99 +++++++++++++++++++++++-----------------
 1 file changed, 58 insertions(+), 41 deletions(-)

diff --git a/test/py/u_boot_console_base.py b/test/py/u_boot_console_base.py
index 384fd53c65..afae07d9cc 100644
--- a/test/py/u_boot_console_base.py
+++ b/test/py/u_boot_console_base.py
@@ -139,8 +139,53 @@ class ConsoleBase(object):
             self.p.close()
         self.logstream.close()
 
+    def wait_for_boot_prompt(self):
+        """Wait for the boot up until command prompt. This is for internal use only.
+        """
+        try:
+            bcfg = self.config.buildconfig
+            config_spl = bcfg.get('config_spl', 'n') == 'y'
+            config_spl_serial = bcfg.get('config_spl_serial', 'n') == 'y'
+            env_spl_skipped = self.config.env.get('env__spl_skipped', False)
+            env_spl2_skipped = self.config.env.get('env__spl2_skipped', True)
+
+            if config_spl and config_spl_serial and not env_spl_skipped:
+                m = self.p.expect([pattern_u_boot_spl_signon] +
+                                  self.bad_patterns)
+                if m != 0:
+                    raise Exception('Bad pattern found on SPL console: ' +
+                                    self.bad_pattern_ids[m - 1])
+            if not env_spl2_skipped:
+                m = self.p.expect([pattern_u_boot_spl2_signon] +
+                                  self.bad_patterns)
+                if m != 0:
+                    raise Exception('Bad pattern found on SPL2 console: ' +
+                                    self.bad_pattern_ids[m - 1])
+            m = self.p.expect([pattern_u_boot_main_signon] + self.bad_patterns)
+            if m != 0:
+                raise Exception('Bad pattern found on console: ' +
+                                self.bad_pattern_ids[m - 1])
+            self.u_boot_version_string = self.p.after
+            while True:
+                m = self.p.expect([self.prompt_compiled,
+                    pattern_stop_autoboot_prompt] + self.bad_patterns)
+                if m == 0:
+                    break
+                if m == 1:
+                    self.p.send(' ')
+                    continue
+                raise Exception('Bad pattern found on console: ' +
+                                self.bad_pattern_ids[m - 2])
+
+        except Exception as ex:
+            self.log.error(str(ex))
+            self.cleanup_spawn()
+            raise
+        finally:
+            self.log.timestamp()
+
     def run_command(self, cmd, wait_for_echo=True, send_nl=True,
-            wait_for_prompt=True):
+            wait_for_prompt=True, wait_for_reboot=False):
         """Execute a command via the U-Boot console.
 
         The command is always sent to U-Boot.
@@ -168,6 +213,9 @@ class ConsoleBase(object):
             wait_for_prompt: Boolean indicating whether to wait for the
                 command prompt to be sent by U-Boot. This typically occurs
                 immediately after the command has been executed.
+            wait_for_reboot: Boolean indication whether to wait for the
+                reboot U-Boot. If this sets True, wait_for_prompt must also
+                be True.
 
         Returns:
             If wait_for_prompt == False:
@@ -202,11 +250,14 @@ class ConsoleBase(object):
                                     self.bad_pattern_ids[m - 1])
             if not wait_for_prompt:
                 return
-            m = self.p.expect([self.prompt_compiled] + self.bad_patterns)
-            if m != 0:
-                self.at_prompt = False
-                raise Exception('Bad pattern found on console: ' +
-                                self.bad_pattern_ids[m - 1])
+            if wait_for_reboot:
+                self.wait_for_boot_prompt()
+            else:
+                m = self.p.expect([self.prompt_compiled] + self.bad_patterns)
+                if m != 0:
+                    self.at_prompt = False
+                    raise Exception('Bad pattern found on console: ' +
+                                    self.bad_pattern_ids[m - 1])
             self.at_prompt = True
             self.at_prompt_logevt = self.logstream.logfile.cur_evt
             # Only strip \r\n; space/TAB might be significant if testing
@@ -349,41 +400,7 @@ class ConsoleBase(object):
             if not self.config.gdbserver:
                 self.p.timeout = 30000
             self.p.logfile_read = self.logstream
-            bcfg = self.config.buildconfig
-            config_spl = bcfg.get('config_spl', 'n') == 'y'
-            config_spl_serial = bcfg.get('config_spl_serial',
-                                                 'n') == 'y'
-            env_spl_skipped = self.config.env.get('env__spl_skipped',
-                                                  False)
-            env_spl2_skipped = self.config.env.get('env__spl2_skipped',
-                                                  True)
-            if config_spl and config_spl_serial and not env_spl_skipped:
-                m = self.p.expect([pattern_u_boot_spl_signon] +
-                                  self.bad_patterns)
-                if m != 0:
-                    raise Exception('Bad pattern found on SPL console: ' +
-                                    self.bad_pattern_ids[m - 1])
-            if not env_spl2_skipped:
-                m = self.p.expect([pattern_u_boot_spl2_signon] +
-                                  self.bad_patterns)
-                if m != 0:
-                    raise Exception('Bad pattern found on SPL2 console: ' +
-                                    self.bad_pattern_ids[m - 1])
-            m = self.p.expect([pattern_u_boot_main_signon] + self.bad_patterns)
-            if m != 0:
-                raise Exception('Bad pattern found on console: ' +
-                                self.bad_pattern_ids[m - 1])
-            self.u_boot_version_string = self.p.after
-            while True:
-                m = self.p.expect([self.prompt_compiled,
-                    pattern_stop_autoboot_prompt] + self.bad_patterns)
-                if m == 0:
-                    break
-                if m == 1:
-                    self.p.send(' ')
-                    continue
-                raise Exception('Bad pattern found on console: ' +
-                                self.bad_pattern_ids[m - 2])
+            self.wait_for_boot_prompt()
             self.at_prompt = True
             self.at_prompt_logevt = self.logstream.logfile.cur_evt
         except Exception as ex: