CS168 proj2 test runner w/ filtering

test_suite.py

#!/usr/bin/env python2.7
"""
Test runner for dv_router.py and learning_switch.py.

Add your own tests by creating new files in tests/ and updating main
below.

"""

from __future__ import print_function

import os
import select
import subprocess
import time

import argparse

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-f",
        "--filter",
        help="only run tests whose name contains this string")
    args = parser.parse_args()

    t = TestSuite()

    # Add or remove learning switch tests here.
    learning_switch_tests = [
        ('tests.test_learning', [])
    ]

    # Add or remove dv router tests here:
    # (test_name, [<extra args>])
    dv_router_tests = [
        ('tests.test_simple', []),
        ('tests.test_failure', []),
        ('tests.test_initialize_neighbor', []),
        ('tests.test_no_hairpin', []),
        ('tests.test_link_weights', []),
        ('tests.test_expire_routes', []),
        ('tests.test_route_poisoning', ['--poison-mode']),
        ('tests.test_count_to_infinity', ['--poison-mode']),
        ('tests.test_count_to_infinity_v2', ['--poison-mode']),
        ('tests.test_multiple_routers', []),
        ('tests.test_fallback_direct', []),
        ('tests.test_equal_cost_paths', []),
        ('tests.test_topo_candy', ['--poison-mode']),
        ('tests.test_topo_rand', ['--poison-mode'])
    ]
    if args.filter:
        print("Running tests with '{0}' in the name".format(args.filter))
        learning_switch_tests = list(
            filter(lambda t: args.filter in t[0], learning_switch_tests))
        dv_router_tests = list(
            filter(lambda t: args.filter in t[0], dv_router_tests))

    for test in dv_router_tests:
        test_name, test_args = test
        t.test('dv_router', test_name, extra_args=test_args)

    # Add your own tests here.

    t.finish()


GREEN = '\033[92m'
RED = '\033[91m'
CLEAR = '\033[0m'


class TestSuite:
    num_passed = 0
    num_failed = 0
    failed = []

    def test(self, router, test_name, extra_args=None):
        cmd = ['python', 'simulator.py', '--no-interactive', '--virtual-time',
               '--default-switch-type=%s' % router]
        if extra_args:
            cmd += extra_args
        cmd += [test_name]
        interactive_cmd = [
            arg for arg in cmd
            if arg not in [
                '--no-interactive', '--virtual-time'
            ]
        ]
        r = subprocess.call(cmd)
        if r is False:
            self.fail(router, test_name, interactive_cmd, 'Timed out')
        elif r is None:
            self.fail(router, test_name, interactive_cmd, 'Could not run')
        elif r == 0:
            print(test_name, r)
            self.succeed(router, test_name)
        else:
            self.fail(router, test_name, interactive_cmd)

    def succeed(self, router, testname):
        print('%s*** %s: %s passed ***%s' % (GREEN, router, testname, CLEAR))
        self.num_passed += 1

    def fail(self, router, testname, cmd, message=None):
        if message:
            print('%s*** %s: %s failed: %s ***%s' % (RED, router, testname,
                                                     message, CLEAR))
        else:
            print('%s*** %s: %s failed ***%s' % (RED, router, testname, CLEAR))
        print('%sCommand: %s%s' % (RED, ' '.join(cmd), CLEAR))
        self.num_failed += 1
        self.failed.append(testname)

    def finish(self):
        if self.num_failed == 0:
            print('%sAll tests passed.%s' % (GREEN, CLEAR))
        else:
            print('Tests: %d passed, %s%d failed%s.' %
                  (self.num_passed, RED, self.num_failed, CLEAR))
            print('Failed: %s%s%s' % (RED, ', '.join(self.failed), CLEAR))


if __name__ == '__main__':
    main()