test/py/u_boot_spawn.py

   1 # SPDX-License-Identifier: GPL-2.0
   2 # Copyright (c) 2015-2016, NVIDIA CORPORATION. All rights reserved.
   3
   4 # Logic to spawn a sub-process and interact with its stdio.
   5
   6 import os
   7 import re
   8 import pty
   9 import signal
  10 import select
  11 import time
  12
  13 class Timeout(Exception):
  14     """An exception sub-class that indicates that a timeout occurred."""
  15     pass
  16
  17 class Spawn(object):
  18     """Represents the stdio of a freshly created sub-process. Commands may be
  19     sent to the process, and responses waited for.
  20
  21     Members:
  22         output: accumulated output from expect()
  23     """
  24
  25     def __init__(self, args, cwd=None):
  26         """Spawn (fork/exec) the sub-process.
  27
  28         Args:
  29             args: array of processs arguments. argv[0] is the command to
  30               execute.
  31             cwd: the directory to run the process in, or None for no change.
  32
  33         Returns:
  34             Nothing.
  35         """
  36
  37         self.waited = False
  38         self.buf = ''
  39         self.output = ''
  40         self.logfile_read = None
  41         self.before = ''
  42         self.after = ''
  43         self.timeout = None
  44         # http://stackoverflow.com/questions/7857352/python-regex-to-match-vt100-escape-sequences
  45         # Note that re.I doesn't seem to work with this regex (or perhaps the
  46         # version of Python in Ubuntu 14.04), hence the inclusion of a-z inside
  47         # [] instead.
  48         self.re_vt100 = re.compile('(\x1b\[|\x9b)[^@-_a-z]*[@-_a-z]|\x1b[@-_a-z]')
  49
  50         (self.pid, self.fd) = pty.fork()
  51         if self.pid == 0:
  52             try:
  53                 # For some reason, SIGHUP is set to SIG_IGN at this point when
  54                 # run under "go" (www.go.cd). Perhaps this happens under any
  55                 # background (non-interactive) system?
  56                 signal.signal(signal.SIGHUP, signal.SIG_DFL)
  57                 if cwd:
  58                     os.chdir(cwd)
  59                 os.execvp(args[0], args)
  60             except:
  61                 print('CHILD EXECEPTION:')
  62                 import traceback
  63                 traceback.print_exc()
  64             finally:
  65                 os._exit(255)
  66
  67         try:
  68             self.poll = select.poll()
  69             self.poll.register(self.fd, select.POLLIN | select.POLLPRI | select.POLLERR | select.POLLHUP | select.POLLNVAL)
  70         except:
  71             self.close()
  72             raise
  73
  74     def kill(self, sig):
  75         """Send unix signal "sig" to the child process.
  76
  77         Args:
  78             sig: The signal number to send.
  79
  80         Returns:
  81             Nothing.
  82         """
  83
  84         os.kill(self.pid, sig)
  85
  86     def isalive(self):
  87         """Determine whether the child process is still running.
  88
  89         Args:
  90             None.
  91
  92         Returns:
  93             Boolean indicating whether process is alive.
  94         """
  95
  96         if self.waited:
  97             return False
  98
  99         w = os.waitpid(self.pid, os.WNOHANG)
 100         if w[0] == 0:
 101             return True
 102
 103         self.waited = True
 104         return False
 105
 106     def send(self, data):
 107         """Send data to the sub-process's stdin.
 108
 109         Args:
 110             data: The data to send to the process.
 111
 112         Returns:
 113             Nothing.
 114         """
 115
 116         os.write(self.fd, data)
 117
 118     def expect(self, patterns):
 119         """Wait for the sub-process to emit specific data.
 120
 121         This function waits for the process to emit one pattern from the
 122         supplied list of patterns, or for a timeout to occur.
 123
 124         Args:
 125             patterns: A list of strings or regex objects that we expect to
 126                 see in the sub-process' stdout.
 127
 128         Returns:
 129             The index within the patterns array of the pattern the process
 130             emitted.
 131
 132         Notable exceptions:
 133             Timeout, if the process did not emit any of the patterns within
 134             the expected time.
 135         """
 136
 137         for pi in range(len(patterns)):
 138             if type(patterns[pi]) == type(''):
 139                 patterns[pi] = re.compile(patterns[pi])
 140
 141         tstart_s = time.time()
 142         try:
 143             while True:
 144                 earliest_m = None
 145                 earliest_pi = None
 146                 for pi in range(len(patterns)):
 147                     pattern = patterns[pi]
 148                     m = pattern.search(self.buf)
 149                     if not m:
 150                         continue
 151                     if earliest_m and m.start() >= earliest_m.start():
 152                         continue
 153                     earliest_m = m
 154                     earliest_pi = pi
 155                 if earliest_m:
 156                     pos = earliest_m.start()
 157                     posafter = earliest_m.end()
 158                     self.before = self.buf[:pos]
 159                     self.after = self.buf[pos:posafter]
 160                     self.output += self.buf[:posafter]
 161                     self.buf = self.buf[posafter:]
 162                     return earliest_pi
 163                 tnow_s = time.time()
 164                 if self.timeout:
 165                     tdelta_ms = (tnow_s - tstart_s) * 1000
 166                     poll_maxwait = self.timeout - tdelta_ms
 167                     if tdelta_ms > self.timeout:
 168                         raise Timeout()
 169                 else:
 170                     poll_maxwait = None
 171                 events = self.poll.poll(poll_maxwait)
 172                 if not events:
 173                     raise Timeout()
 174                 c = os.read(self.fd, 1024)
 175                 if not c:
 176                     raise EOFError()
 177                 if self.logfile_read:
 178                     self.logfile_read.write(c)
 179                 self.buf += c
 180                 # count=0 is supposed to be the default, which indicates
 181                 # unlimited substitutions, but in practice the version of
 182                 # Python in Ubuntu 14.04 appears to default to count=2!
 183                 self.buf = self.re_vt100.sub('', self.buf, count=1000000)
 184         finally:
 185             if self.logfile_read:
 186                 self.logfile_read.flush()
 187
 188     def close(self):
 189         """Close the stdio connection to the sub-process.
 190
 191         This also waits a reasonable time for the sub-process to stop running.
 192
 193         Args:
 194             None.
 195
 196         Returns:
 197             Nothing.
 198         """
 199
 200         os.close(self.fd)
 201         for i in range(100):
 202             if not self.isalive():
 203                 break
 204             time.sleep(0.1)
 205
 206     def get_expect_output(self):
 207         """Return the output read by expect()
 208
 209         Returns:
 210             The output processed by expect(), as a string.
 211         """
 212         return self.output