testrunner

#!/usr/bin/env python

from __future__ import print_function

import json
import os
import shlex
import socket
import sys
import time
import xml.etree.ElementTree as ET

from argparse import ArgumentParser, Action
from collections import defaultdict
from fcntl import fcntl, F_GETFL, F_SETFL
from fnmatch import fnmatch
from multiprocessing import cpu_count
from select import select
from subprocess import Popen, PIPE

try:
    # Available since python 3.3
    from shlex import quote
except ImportError:
    from pipes import quote

class Colors:
    GRAY = "\x1b[1;30m"
    RED = "\x1b[1;31m"
    GREEN = "\x1b[1;32m"
    DEFAULT = "\x1b[0;39m"
    REVERSE = "\x1b[7m"


def get_dir(path):
    d = os.path.dirname(path)
    if d == '':
        return '.'
    return d


class TestCase:

    def __init__(self, name=None, description=None, tags=[], **kwargs):
        self.name = name
        self.description = description
        self.tags = tags
        self.output_file_path = None
        self.crashed = False

    def __repr__(self):
        return self.name


class TestModule:

    def __init__(self, module_path, idx):
        self.module_path = module_path
        self.name = "Unnamed module (%s)" % module_path
        self.idx = idx
        self.finished = False
        self.failed = False
        self.tests = []
        self.read_buffer = b''

    def list(self):
        args = [os.path.abspath(self.module_path), '--list']
        self.proc = Popen(args, stdout=PIPE, cwd=get_dir(self.module_path))

    def run(self, test_indices, wrapper):
        args = []
        args.extend([os.path.abspath(self.module_path), '--run'])
        args.extend(map(str, test_indices))
        args = wrapper.get_args(args)
        self.proc = Popen(args, stdout=PIPE, cwd=get_dir(self.module_path))
        flags = fcntl(self.fileno(), F_GETFL)
        fcntl(self.fileno(), F_SETFL, flags | os.O_NONBLOCK)

    def read_events(self):
        # Check the status of the process
        self.proc.poll()

        self.read_buffer += self.proc.stdout.read()
        if b'\n' in self.read_buffer:
            lines = self.read_buffer.split(b'\n')
            # Yield all pending events
            for line in lines[:-1]:
                try:
                    event = json.loads(line.decode())
                except ValueError:
                    event = {
                        'event': 'testcase_error',
                        'message': "Failed to parse test case output",
                        'crash': False
                    }
                yield event
            self.read_buffer = lines[-1]

        # Handle process completion
        if self.proc.returncode is not None:
            self.finished = True
            if self.proc.returncode != 0:
                self.failed = True

    def reset_status(self):
        self.finished = False
        self.failed = False

    def fileno(self):
        return self.proc.stdout.fileno()

    def read(self, size):
        return self.proc.stdout.read(size)


class Observer:

    def call(self, module, event):
        handler = getattr(self, 'handle_' + event['event'], None)
        if handler:
            handler(module, event)


class EventEmitter(object):

    def __init__(self):
        self.observers = []

    def register(self, observer):
        self.observers.append(observer)

    def deregister(self, observer):
        self.observers.remove(observer)

    def emit(self, module, event):
        for observer in self.observers:
            observer.call(module, event)


class BufferedEventEmitter(EventEmitter):

    def __init__(self):
        super(BufferedEventEmitter, self).__init__()
        self.current_module = None
        self.buffered_events = defaultdict(list)
        self.finished_modules = []

    def call(self, module, event):
        if event['event'] == 'module_start':
            self.handle_module_start(module)

        if module == self.current_module:
            self.emit(module, event)
        else:
            self.buffered_events[module].append(event)

        if event['event'] == 'module_end':
            self.handle_module_end(module)

    def handle_module_start(self, module):
        if not self.current_module:
            self.current_module = module

    def handle_module_end(self, module):
        if module == self.current_module:
            while self.finished_modules:
                next_module = self.finished_modules.pop()
                for e in self.buffered_events.pop(next_module):
                    self.emit(next_module, e)
            if self.buffered_events:
                next_module, events = self.buffered_events.popitem()
                self.current_module = next_module
                for e in events:
                    self.emit(next_module, e)
            else:
                self.current_module = None
        else:
            self.finished_modules.append(module)


class Runner(EventEmitter):

    def __init__(self, module_paths, jobs):
        super(Runner, self).__init__()
        self.modules = [TestModule(t, i) for i, t in enumerate(module_paths)]
        self.simultaneous_jobs = jobs

    def list_modules(self):
        self.reset_modules()
        pending_jobs = self.modules[:]
        running_jobs = []
        while pending_jobs or running_jobs:
            while len(running_jobs) < self.simultaneous_jobs and pending_jobs:
                job = pending_jobs.pop()
                job.list()
                running_jobs.append(job)
            job = self.handle_events(running_jobs)
            running_jobs.remove(job)

    def run_modules(self, tests_to_run, wrapperclass, args):
        self.reset_modules()
        pending_jobs = tests_to_run[:]
        running_jobs = []
        while pending_jobs or running_jobs:
            while len(running_jobs) < self.simultaneous_jobs and pending_jobs:
                job, indices = pending_jobs.pop()
                wrapper = wrapperclass(job, args)
                job.run(indices, wrapper)
                self.emit(job, {
                    'event': 'module_start',
                    'message': wrapper.get_message(),
                })
                running_jobs.append(job)
            job = self.handle_events(running_jobs)
            running_jobs.remove(job)

    def reset_modules(self):
        for m in self.modules:
            m.reset_status()

    def handle_events(self, modules):
        periodic_next = time.time() + 0.1
        while True:
            # Attempt to read from all running modules
            rs, _, _ = select(modules, [], [], 0.1)
            if time.time() > periodic_next:
                periodic_next += 0.1
                for m in modules:
                    self.emit(m, {
                        'event': 'periodic'
                    })
            for r in rs:
                # Handle all pending events
                for event in r.read_events():
                    self.emit(r, event)
                # Handle module completion
                if r.finished:
                    if r.failed:
                        self.emit(r, {
                            'event': 'module_crash',
                            'message': "Test module crashed",
                        })
                    self.emit(r, {
                        'event': 'module_end',
                    })
                    return r


class TestModuleCollector(Observer):

    def handle_testcase_list(self, module, event):
        module.tests.append(TestCase(**event))

    def handle_module_list(self, module, event):
        module.name = event['name']


def hat_escape(data):
    r = []
    for c in data:
        if ord(c) < ord(' '):
            r.append("{colors.REVERSE}^{c}{colors.DEFAULT}"
                     .format(c=chr(ord("@") + ord(c)),
                             colors=Colors))
        else:
            r.append(c)
    return "".join(r)


def print_output_lines(lines):
    for line in lines:
        if line.endswith(b"\n"):
            line = line[:-1]
        line = line.decode("utf-8", "replace")
        print("           {colors.GRAY}>{colors.DEFAULT} {line}"
              .format(line=hat_escape(line), colors=Colors))


def print_output_file(path):
    with open(path, "rb") as f:
        print_output_lines(f)


class TestOutputPrinter(Observer):
    def __init__(self):
        self.modulestate = {}

    def handle_testcase_start(self, module, event):
        prefix = module.tests[event['index']].description
        output_file = open(event['output'], 'rb')
        self.modulestate[module] = (prefix, output_file)
        self._emit_output(*self.modulestate[module])

    def handle_periodic(self, module, event):
        if module in self.modulestate:
            self._emit_output(*self.modulestate[module])

    def handle_testcase_end(self, module, event):
        self._emit_output(*self.modulestate[module])
        prefix, output_file = self.modulestate[module]
        del self.modulestate[module]
        output_file.close()

    def _emit_output(self, prefix, output_file):
        pos = output_file.tell()
        # Seek to current position clears EOF state in case there is more data
        output_file.seek(pos)
        data = output_file.read()
        if data:
            if pos == 0:
                print("    [      ] " + prefix)
            print_output_lines(data.split(b"\n"))


class TestEmitter(Observer):

    def __init__(self, show_output):
        self.current_test = None
        self.show_output = show_output
        self.nontestoutput = None

    def handle_module_start(self, module, event):
        print("  {name}".format(name=module.name))
        for line in event['message']:
            print("    > {}".format(line))

    def handle_setup_start(self, module, event):
        assert self.nontestoutput is None
        self.nontestoutput = event['output']

    def handle_setup_end(self, module, event):
        assert self.nontestoutput is not None
        self.nontestoutput = None

    def handle_teardown_start(self, module, event):
        assert self.nontestoutput is None
        self.nontestoutput = event['output']

    def handle_teardown_end(self, module, event):
        assert self.nontestoutput is not None
        self.nontestoutput = None

    def handle_testcase_start(self, module, event):
        self.current_test = module.tests[event['index']]
        self.current_test.output_file_path = event['output']

    def handle_testcase_end(self, module, event):
        self.print_testcase(self.current_test, event)
        self.current_test = None

    def handle_testcase_error(self, module, event):
        # When a test case error event is generated by the test itself,
        # we will get an additional one created by the test runner
        # which should we ignored
        if not self.current_test.crashed:
            self.print_testcase(self.current_test, event)
        self.current_test.crashed = True

    def handle_module_crash(self, module, event):
        if self.current_test:
            self.print_testcase(self.current_test, event)
        else:
            print("    [ {colors.RED}FAIL{colors.DEFAULT} ]"
                  .format(colors=Colors))
            print("           ! " + event['message'])
            if self.nontestoutput is not None:
                print_output_file(self.nontestoutput)
                self.nontestoutput = None

        self.current_test = None

    def print_testcase(self, test, event):
        event_type = event['event']
        success = event['success'] if event_type == 'testcase_end' else False
        result = (Colors.GREEN + "PASS") if success else (Colors.RED + "FAIL")
        print(
            "    [ {result}{colors.DEFAULT} ] {desc}"
            .format(desc=test.description, result=result, colors=Colors), end=''
        )
        siz = os.path.getsize(test.output_file_path)
        if siz > 0:
            print(" {colors.GRAY}({siz} bytes of output generated){colors.DEFAULT}"
                  .format(siz=siz, colors=Colors), end='')

        if event_type == 'testcase_end':
            print(
                " {colors.GRAY}({event[duration]:.3f} s){colors.DEFAULT}"
                .format(event=event, colors=Colors)
            )
            if not success:
                for failure in event['failures']:
                    print("           * " + failure['message'])
                    print("             @ {file}:{line}".format(**failure))
            if event.get('valgrind_errors', 0):
                print("           ! {colors.RED}{vgerrs} valgrind error(s){colors.DEFAULT} reported!"
                      .format(vgerrs=event['valgrind_errors'], colors=Colors))
        elif event_type in ('testcase_error', 'module_crash'):
            print('')
            print("           ! " + event['message'])
            if 'file' in event and 'line' in event:
                print("             @ {file}:{line}".format(**event))
        if not self.show_output and not success:
            print_output_file(test.output_file_path)


class SummaryEmitter(Observer):

    def __init__(self, module_init_failures):
        self.module_counter = 0
        self.module_fail_counter = 0
        self.module_init_fail_counter = module_init_failures
        self.test_counter = 0
        self.test_fail_counter = 0
        self.assert_counter = 0
        self.assert_fail_counter = 0
        self.duration_total = 0
        self.cpu_time_total = 0
        self.has_reported_failing_test = {}
        self.tests_with_valgrind_errors_counter = 0
        self.valgrind_errors_counter = 0
        self.show_valgrind_stats = False

    def handle_module_start(self, module, event):
        self.has_reported_failing_test[module] = False

    def handle_testcase_end(self, module, event):
        self.assert_counter += event['asserts']
        self.assert_fail_counter += len(event['failures'])
        self.test_counter += 1
        self.duration_total += event['duration']
        self.cpu_time_total += event['cpu_time']
        if event['failures']:
            self.test_fail_counter += 1
            self.has_reported_failing_test[module] = True
        if event.get('valgrind_errors', 0):
            self.tests_with_valgrind_errors_counter += 1
            self.valgrind_errors_counter += event['valgrind_errors']

    def handle_testcase_error(self, module, event):
        self.test_counter += 1
        self.test_fail_counter += 1
        if event['crash']:
            self.module_fail_counter += 1

    def handle_module_crash(self, module, event):
        self.module_fail_counter += 1

    def handle_module_end(self, module, event):
        self.module_counter += 1
        if self.has_reported_failing_test[module]:
            self.module_fail_counter += 1

    def is_failure(self):
        return any((self.module_fail_counter > 0,
                    self.valgrind_errors_counter > 0,
                    self.module_init_fail_counter > 0))

    def print_summary(self):
        print("\n  Run summary:\n")
        print((
            "  +----------+------------+------------+------------+\n"
            "  |          |       Pass |       Fail |      Total |\n"
            "  +----------+------------+------------+------------+\n"
            "  | Modules  | {:10} | {:10} | {:10} |\n"
            "  | Tests    | {:10} | {:10} | {:10} |\n"
            "  | Asserts  | {:10} | {:10} | {:10} |\n"
            "  | Elapsed  |            |            | {:9.3f}s |\n"
            "  | CPU time |            |            | {:9.3f}s |\n"
            "  +----------+------------+------------+------------+\n"
        ).format(
            self.module_counter - self.module_fail_counter,
            self.module_fail_counter + self.module_init_fail_counter,
            self.module_counter + self.module_init_fail_counter,
            self.test_counter - self.test_fail_counter,
            self.test_fail_counter,
            self.test_counter,
            self.assert_counter - self.assert_fail_counter,
            self.assert_fail_counter,
            self.assert_counter,
            self.duration_total,
            self.cpu_time_total
        ))
        if self.show_valgrind_stats:
            print((
                "  +--------------------+------------+\n"
                "  | Valgrind           |            |\n"
                "  +--------------------+------------+\n"
                "  |            Errors: | {:10} |\n"
                "  | Tests with errors: | {:10} |\n"
                "  +--------------------+------------+\n"
            ).format(
                self.valgrind_errors_counter,
                self.tests_with_valgrind_errors_counter,
            ))


class XMLEmitter(Observer):

    def __init__(self, xml_path):
        self.xml_path = xml_path
        self.root = ET.Element("testsuites")
        self.current_module = None
        self.current_test = None
        self.current_test_output = None

    def handle_module_start(self, module, event):
        assert self.current_module is None
        self.current_module = ET.SubElement(self.root, "testsuite",
                                            name=module.module_path,
                                            hostname=socket.getfqdn(),
                                            timestamp=time.strftime("%Y-%m-%dT%H:%M:%S"))
        ET.SubElement(self.current_module, "properties")

    def handle_module_end(self, module, event):
        self.current_module = None

    def handle_testcase_start(self, module, event):
        assert self.current_module is not None
        assert self.current_test is None
        self.current_test = ET.SubElement(self.current_module, "testcase",
                                          name=module.tests[event['index']].name,
                                          classname="%s:%s" % (module.module_path, module.tests[event['index']].name))
        self.current_test_output = event['output']

    def add_test_output(self):
        out = ET.SubElement(self.current_test, "system-out")
        with open(self.current_test_output, "rb") as outfile:
            outfile.seek(0, 2)
            siz = outfile.tell()
            if siz > 768:
                outfile.seek(-512, 2)
                prefix = "[truncated (total size was %d bytes)]\n" % siz
            else:
                outfile.seek(0)
                prefix = ""
            data = "".join(("\\x%02x" % ord(c)) if ord(c) < 32 and c != "\n" else c
                           for c in outfile.read().decode("utf-8", "replace"))
            out.text = prefix + data

    def handle_testcase_end(self, module, event):
        assert self.current_test is not None
        self.current_test.attrib["time"] = "%.3f" % event['duration']

        for f in event['failures']:
            ET.SubElement(self.current_test, "failure",
                          message=f['message'], type="assert")

        self.add_test_output()
        self.current_test = None

    def handle_testcase_error(self, module, event):
        ET.SubElement(self.current_test, "error",
                      message=event['message'], type="internal_testcase_error")

        self.add_test_output()
        self.current_test = None

    def handle_module_crash(self, module, event):
        fake_test = ET.SubElement(self.current_module, "testcase",
                                  name="<No Test>",
                                  classname="%s:<No Test>" % module.module_path)
        ET.SubElement(fake_test, "error",
                      message=event['message'], type="crash")
        self.current_test = None

    def write_output(self):
        assert self.root.tag == "testsuites"

        for suite in self.root:
            assert suite.tag == "testsuite"
            cases = errors = failures = 0
            for case in suite:
                if case.tag == "properties":
                    continue
                assert case.tag == "testcase"
                cases += 1
                for s in case:
                    if s.tag == "error":
                        errors += 1
                    elif s.tag == "failure":
                        failures += 1
            suite.attrib["tests"] = str(cases)
            suite.attrib["errors"] = str(errors)
            suite.attrib["failures"] = str(failures)

        doc = ET.ElementTree(xml_emitter.root)
        with open(self.xml_path, "wb") as out:
            doc.write(out, encoding="UTF-8", xml_declaration=True)


class TestCleaner(Observer):

    def __init__(self):
        self.output_file_path = None

    def handle_testcase_start(self, module, event):
        assert self.output_file_path is None
        self.output_file_path = event['output']

    def handle_testcase_end(self, module, event):
        self.clean_output()

    def handle_module_crash(self, module, event):
        self.clean_output()

    def handle_testcase_error(self, module, event):
        self.clean_output()

    def handle_setup_start(self, module, event):
        assert self.output_file_path is None
        self.output_file_path = event['output']

    def handle_setup_end(self, module, event):
        self.clean_output()

    def handle_teardown_start(self, module, event):
        assert self.output_file_path is None
        self.output_file_path = event['output']

    def handle_teardown_end(self, module, event):
        self.clean_output()

    def clean_output(self):
        if self.output_file_path:
            os.unlink(self.output_file_path)
            self.output_file_path = None


class FilterAction(Action):
    def __call__(self, parser, namespace, values, option_string):
        items = getattr(namespace, self.dest)
        if items is None:
            items = []
        items.append((values, namespace.exclude))
        setattr(namespace, self.dest, items)
        namespace.exclude = False


class Wrapper(object):
    def __init__(self, module, mainargs):
        self.module = module
        self.mainargs = mainargs

    def get_args(self, args):
        return args

    def get_message(self):
        return []


class Valgrind(Wrapper):
    log_path = property(lambda self:
                        "valgrind.{}.log".format(os.path.basename(self.module.module_path)))
    extra_opts = property(lambda self: self.mainargs.valgrind_opt)

    def get_args(self, args):
        outargs = ['valgrind', "--log-file={}".format(self.log_path)]
        outargs += self.extra_opts
        outargs += args
        return outargs

    def get_message(self):
        mesg = ["VALGRIND MODE ENABLED",
                "",
                "The test process will be run under Valgrind. Any errors found will be logged to:",
                "\t{}".format(self.log_path),
                ""]
        if self.extra_opts:
            mesg += ["Extra options added:",
                     "\t{}".format(" ".join(map(quote, self.extra_opts))),
                     ""]
        return mesg


class GDBServer(Wrapper):
    comm = property(lambda self:
                    os.getenv('SCU_GDBSERVER_COMM', '127.0.0.1:9999'))

    def get_args(self, args):
        return ['gdbserver', self.comm] + args

    def get_message(self):
        return ["DEBUG MODE ENABLED",
                "",
                "The test process will be run under gdbserver",
                "To attach to it run the following command in gdb:",
                " (gdb) target remote {}".format(self.comm),
                "Then to start it:",
                " (gdb) continue"]


if __name__ == '__main__':
    # Parse arguments

    show_output_default = os.getenv("SCU_SHOW_OUTPUT", "") not in ("", "0")
    valgrind_opts_default = shlex.split(os.getenv("SCU_VALGRIND_OPTS", ""))

    parser = ArgumentParser(description="Runs SCU test modules")
    parser.add_argument('module', nargs='+', help="path to test module")
    parser.add_argument('--name', action=FilterAction, default=[], help="filter test cases by name")
    parser.add_argument('--tag', action=FilterAction, default=[], help="filter test cases by tag")
    parser.add_argument('--exclude', action='store_true', help="negates the following filter option")
    parser.add_argument('-g', '--gdb', action='store_true', help="gdb compatibility mode")
    parser.add_argument('-v', '--valgrind', action='store_true', help="valgrind compatibility mode")
    parser.add_argument('--valgrind-opt', action='append', default=valgrind_opts_default,
                        help="extra option to pass to valgrind")
    parser.add_argument('-j', '--jobs', default=cpu_count(), type=int, help="number of jobs to run simultaneously")
    parser.add_argument('--show-output', action='store_true', default=show_output_default, help="show test stdout/err")
    parser.add_argument('--xml', help="store results to xml file (junitxml)")
    args = parser.parse_args()

    # Create runner
    runner = Runner(args.module, args.jobs)

    # List all tests
    collector = TestModuleCollector()
    runner.register(collector)
    runner.list_modules()
    runner.deregister(collector)

    module_init_failures = 0
    for m in runner.modules:
        if m.failed:
            module_init_failures += 1
            print("  {}".format(m.module_path))
            print("    > {colors.RED}Module failed to initialize{colors.DEFAULT}"
                  .format(colors=Colors))
        elif not m.tests:
            print("  {name}".format(name=m.name))
            print("    > Module does not contain any tests"
                  .format(colors=Colors))

    # Filter tests
    tests_to_run = []
    for m in runner.modules:
        indices = set(range(len(m.tests)))
        for p, exclude in args.name:
            s = set(i for i, t in enumerate(m.tests) if fnmatch(t.name, p))
            if exclude:
                s = indices - s
            indices = indices & s
        for p, exclude in args.tag:
            s = set(i for i, t in enumerate(m.tests) if p in t.tags)
            if exclude:
                s = indices - s
            indices = indices & s
        if indices:
            tests_to_run.append((m, indices))

    # Set up observers
    buffered_emitter = BufferedEventEmitter()
    if args.show_output:
        buffered_emitter.register(TestOutputPrinter())
    buffered_emitter.register(TestEmitter(args.show_output))
    if args.xml:
        xml_emitter = XMLEmitter(args.xml)
        buffered_emitter.register(xml_emitter)
    else:
        xml_emitter = None
    buffered_emitter.register(TestCleaner())

    summary_emitter = SummaryEmitter(module_init_failures)
    runner.register(buffered_emitter)
    runner.register(summary_emitter)

    if args.gdb:
        wrapperclass = GDBServer
        runner.simultaneous_jobs = 1
    elif args.valgrind:
        summary_emitter.show_valgrind_stats = True
        wrapperclass = Valgrind
    else:
        wrapperclass = Wrapper

    # Run selected tests
    runner.run_modules(tests_to_run, wrapperclass, args)

    # Print summary
    summary_emitter.print_summary()

    if xml_emitter:
        xml_emitter.write_output()

    sys.exit(summary_emitter.is_failure())