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