test/py/u_boot_spawn.py

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