cannam@133: #! /usr/bin/env python
cannam@133: 
cannam@133: # MEGA TEST
cannam@133: #
cannam@133: # usage:  mega-test.py <config>
cannam@133: #
cannam@133: # This runs several tests in parallel and shows progress bars for each, based on a config file.
cannam@133: #
cannam@133: # <config> is a file containing a list of commands to run along with the expected number of lines
cannam@133: # they will output (to stdout and stderr combined), which is how the progress bar is calculated.
cannam@133: # The format of the file is simply one test per line, with the line containing the test name,
cannam@133: # the number of output lines expected, and the test command.  Example:
cannam@133: #
cannam@133: #     mytest 1523 ./my-test --foo bar
cannam@133: #     another 862 ./another-test --baz
cannam@133: #
cannam@133: # Each command is interpreted by `sh -euc`, therefore it is acceptable to use environment
cannam@133: # variables and other shell syntax.
cannam@133: #
cannam@133: # After all tests complete, the config file will be rewritten to update the line counts to the
cannam@133: # actual number of lines seen for all passing tests (failing tests are not updated).
cannam@133: 
cannam@133: import sys
cannam@133: import re
cannam@133: import os
cannam@133: from errno import EAGAIN
cannam@133: from fcntl import fcntl, F_GETFL, F_SETFL
cannam@133: from select import poll, POLLIN, POLLHUP
cannam@133: from subprocess import Popen, PIPE, STDOUT
cannam@133: 
cannam@133: CONFIG_LINE = re.compile("^([^ ]+) +([0-9]+) +(.*)$")
cannam@133: 
cannam@133: if len(sys.argv) != 2:
cannam@133:   sys.stderr.write("Wrong number of arguments.\n");
cannam@133:   sys.exit(1)
cannam@133: 
cannam@133: if not os.access("/tmp/test-output", os.F_OK):
cannam@133:   os.mkdir("/tmp/test-output")
cannam@133: 
cannam@133: config = open(sys.argv[1], 'r')
cannam@133: 
cannam@133: tests = []
cannam@133: 
cannam@133: class Test:
cannam@133:   def __init__(self, name, command, lines):
cannam@133:     self.name = name
cannam@133:     self.command = command
cannam@133:     self.lines = lines
cannam@133:     self.count = 0
cannam@133:     self.done = False
cannam@133: 
cannam@133:   def start(self, poller):
cannam@133:     self.proc = Popen(["sh", "-euc", test.command], stdin=dev_null, stdout=PIPE, stderr=STDOUT)
cannam@133:     fd = self.proc.stdout.fileno()
cannam@133:     flags = fcntl(fd, F_GETFL)
cannam@133:     fcntl(fd, F_SETFL, flags | os.O_NONBLOCK)
cannam@133:     poller.register(self.proc.stdout, POLLIN)
cannam@133:     self.log = open("/tmp/test-output/" + self.name + ".log", "w")
cannam@133: 
cannam@133:   def update(self):
cannam@133:     try:
cannam@133:       while True:
cannam@133:         text = self.proc.stdout.read()
cannam@133:         if text == "":
cannam@133:           self.proc.wait()
cannam@133:           self.done = True
cannam@133:           self.log.close()
cannam@133:           return True
cannam@133:         self.count += text.count("\n")
cannam@133:         self.log.write(text)
cannam@133:     except IOError as e:
cannam@133:       if e.errno == EAGAIN:
cannam@133:         return False
cannam@133:       raise
cannam@133: 
cannam@133:   def print_bar(self):
cannam@133:     percent = self.count * 100 / self.lines
cannam@133:     status = "(%3d%%)" % percent
cannam@133: 
cannam@133:     color_on = ""
cannam@133:     color_off = ""
cannam@133: 
cannam@133:     if self.done:
cannam@133:       if self.proc.returncode == 0:
cannam@133:         color_on = "\033[0;32m"
cannam@133:         status = "PASS"
cannam@133:       else:
cannam@133:         color_on = "\033[0;31m"
cannam@133:         status = "FAIL: /tmp/test-output/%s.log" % self.name
cannam@133:       color_off = "\033[0m"
cannam@133: 
cannam@133:     print "%s%-16s |%-25s| %6d/%6d %s%s    " % (
cannam@133:         color_on, self.name, '=' * min(percent / 4, 25), self.count, self.lines, status, color_off)
cannam@133: 
cannam@133:   def passed(self):
cannam@133:     return self.proc.returncode == 0
cannam@133: 
cannam@133: for line in config:
cannam@133:   if len(line) > 0 and not line.startswith("#"):
cannam@133:     match = CONFIG_LINE.match(line)
cannam@133:     if not match:
cannam@133:       sys.stderr.write("Invalid config syntax: %s\n" % line);
cannam@133:       sys.exit(1)
cannam@133:     test = Test(match.group(1), match.group(3), int(match.group(2)))
cannam@133:     tests.append(test)
cannam@133: 
cannam@133: config.close()
cannam@133: 
cannam@133: dev_null = open("/dev/null", "rw")
cannam@133: poller = poll()
cannam@133: fd_map = {}
cannam@133: 
cannam@133: for test in tests:
cannam@133:   test.start(poller)
cannam@133:   fd_map[test.proc.stdout.fileno()] = test
cannam@133: 
cannam@133: active_count = len(tests)
cannam@133: 
cannam@133: def print_bars():
cannam@133:   for test in tests:
cannam@133:     test.print_bar()
cannam@133: 
cannam@133: print_bars()
cannam@133: 
cannam@133: while active_count > 0:
cannam@133:   for (fd, event) in poller.poll():
cannam@133:     if fd_map[fd].update():
cannam@133:       active_count -= 1
cannam@133:       poller.unregister(fd)
cannam@133:   sys.stdout.write("\033[%dA\r" % len(tests))
cannam@133:   print_bars()
cannam@133: 
cannam@133: new_config = open(sys.argv[1], "w")
cannam@133: for test in tests:
cannam@133:   if test.passed():
cannam@133:     new_config.write("%-16s %6d %s\n" % (test.name, test.count, test.command))
cannam@133:   else:
cannam@133:     new_config.write("%-16s %6d %s\n" % (test.name, test.lines, test.command))
cannam@133: 
cannam@133: for test in tests:
cannam@133:   if not test.passed():
cannam@133:     sys.exit(1)
cannam@133: 
cannam@133: sys.exit(0)