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
  23     def __init__(self, args):
  24         """Spawn (fork/exec) the sub-process.
  25
  26         Args:
  27             args: array of processs arguments. argv[0] is the command to execute.
  28
  29         Returns:
  30             Nothing.
  31         """
  32
  33         self.waited = False
  34         self.buf = ''
  35         self.logfile_read = None
  36         self.before = ''
  37         self.after = ''
  38         self.timeout = None
  39
  40         (self.pid, self.fd) = pty.fork()
  41         if self.pid == 0:
  42             try:
  43                 # For some reason, SIGHUP is set to SIG_IGN at this point when
  44                 # run under "go" (www.go.cd). Perhaps this happens under any
  45                 # background (non-interactive) system?
  46                 signal.signal(signal.SIGHUP, signal.SIG_DFL)
  47                 os.execvp(args[0], args)
  48             except:
  49                 print 'CHILD EXECEPTION:'
  50                 import traceback
  51                 traceback.print_exc()
  52             finally:
  53                 os._exit(255)
  54
  55         self.poll = select.poll()
  56         self.poll.register(self.fd, select.POLLIN | select.POLLPRI | select.POLLERR | select.POLLHUP | select.POLLNVAL)
  57
  58     def kill(self, sig):
  59         """Send unix signal "sig" to the child process.
  60
  61         Args:
  62             sig: The signal number to send.
  63
  64         Returns:
  65             Nothing.
  66         """
  67
  68         os.kill(self.pid, sig)
  69
  70     def isalive(self):
  71         """Determine whether the child process is still running.
  72
  73         Args:
  74             None.
  75
  76         Returns:
  77             Boolean indicating whether process is alive.
  78         """
  79
  80         if self.waited:
  81             return False
  82
  83         w = os.waitpid(self.pid, os.WNOHANG)
  84         if w[0] == 0:
  85             return True
  86
  87         self.waited = True
  88         return False
  89
  90     def send(self, data):
  91         """Send data to the sub-process's stdin.
  92
  93         Args:
  94             data: The data to send to the process.
  95
  96         Returns:
  97             Nothing.
  98         """
  99
 100         os.write(self.fd, data)
 101
 102     def expect(self, patterns):
 103         """Wait for the sub-process to emit specific data.
 104
 105         This function waits for the process to emit one pattern from the
 106         supplied list of patterns, or for a timeout to occur.
 107
 108         Args:
 109             patterns: A list of strings or regex objects that we expect to
 110                 see in the sub-process' stdout.
 111
 112         Returns:
 113             The index within the patterns array of the pattern the process
 114             emitted.
 115
 116         Notable exceptions:
 117             Timeout, if the process did not emit any of the patterns within
 118             the expected time.
 119         """
 120
 121         for pi in xrange(len(patterns)):
 122             if type(patterns[pi]) == type(''):
 123                 patterns[pi] = re.compile(patterns[pi])
 124
 125         tstart_s = time.time()
 126         try:
 127             while True:
 128                 earliest_m = None
 129                 earliest_pi = None
 130                 for pi in xrange(len(patterns)):
 131                     pattern = patterns[pi]
 132                     m = pattern.search(self.buf)
 133                     if not m:
 134                         continue
 135                     if earliest_m and m.start() >= earliest_m.start():
 136                         continue
 137                     earliest_m = m
 138                     earliest_pi = pi
 139                 if earliest_m:
 140                     pos = earliest_m.start()
 141                     posafter = earliest_m.end() + 1
 142                     self.before = self.buf[:pos]
 143                     self.after = self.buf[pos:posafter]
 144                     self.buf = self.buf[posafter:]
 145                     return earliest_pi
 146                 tnow_s = time.time()
 147                 tdelta_ms = (tnow_s - tstart_s) * 1000
 148                 if tdelta_ms > self.timeout:
 149                     raise Timeout()
 150                 events = self.poll.poll(self.timeout - tdelta_ms)
 151                 if not events:
 152                     raise Timeout()
 153                 c = os.read(self.fd, 1024)
 154                 if not c:
 155                     raise EOFError()
 156                 if self.logfile_read:
 157                     self.logfile_read.write(c)
 158                 self.buf += c
 159         finally:
 160             if self.logfile_read:
 161                 self.logfile_read.flush()
 162
 163     def close(self):
 164         """Close the stdio connection to the sub-process.
 165
 166         This also waits a reasonable time for the sub-process to stop running.
 167
 168         Args:
 169             None.
 170
 171         Returns:
 172             Nothing.
 173         """
 174
 175         os.close(self.fd)
 176         for i in xrange(100):
 177             if not self.isalive():
 178                 break
 179             time.sleep(0.1)