import re
import sys
import unittest
from argparse import ArgumentParser


def main():
    ap = ArgumentParser(description="Rockset unittest main")
    ap.add_argument(
        "-b", "--buffer", help="Buffer output during run", action="store_true"
    )
    ap.add_argument(
        "-c", "--catch", help="Catch Ctrl-C during run", action="store_true"
    )
    ap.add_argument(
        "-f", "--failfast", help="Stop on first failure", action="store_true"
    )
    ap.add_argument(
        "--locals", help="Show local variables in tracebacks", action="store_true"
    )
    ap.add_argument("-v", "--verbose", help="Verbose output", action="store_true")
    ap.add_argument("-s", "--start-directory", help="Start directory", default=".")
    ap.add_argument("patterns", help="Test regexes to run", nargs="*")

    args = ap.parse_args()
    patterns = [re.compile(p) for p in args.patterns]

    def filter_tests(tests):
        accepted = []
        for test in tests:
            if isinstance(test, unittest.TestSuite):
                accepted.append(test)
                continue
            test_id = test.id()
            if (not patterns) or any(p.match(test_id) for p in patterns):
                accepted.append(test)

        return unittest.TestSuite(accepted)

    if args.catch:
        unittest.installHandler()

    loader = unittest.TestLoader()
    loader.suiteClass = filter_tests
    suite = loader.discover(args.start_directory)
    runner = unittest.TextTestRunner(
        verbosity=2 if args.verbose else 1, failfast=args.failfast, buffer=args.buffer
    )
    result = runner.run(suite)
    if not result.wasSuccessful():
        sys.exit(1)


if __name__ == "__main__":
    main()
