gray/tests/__main__.py

import argparse
import unittest
import tests


def get_cases(suite):
    '''
    Get a list of TestCase from a TestSuite
    '''
    if hasattr(suite, '__iter__'):
        res = set()
        for x in suite:
            res |= get_cases(x)
        return res

    mod = type(suite).__module__
    cls = type(suite).__name__
    return {f'{mod}.{cls}'}


if __name__ == '__main__':
    # handle custom CLI arguments
    parser = argparse.ArgumentParser(prog='tests.main')
    parser.add_argument('--keep-failed', action='store_true',
                        help='whether to keep output files of failed tests')
    parser.add_argument('--visual', action='store_true',
                        help='whether to show plots, helps with debugging')
    parser.add_argument('--update', action='store_true',
                        help='update the reference results with the'
                             'current ones')
    parser.add_argument('--binary', type=str, default='build/bin/gray',
                        help='the gray binary to be tested')
    parser.add_argument('--list-cases', action='store_true',
                        help='only list the test cases')
    options, other_args = parser.parse_known_args()
    tests.options = options

    if options.list_cases:
        # list all test cases
        suite = unittest.TestLoader().discover('.')
        print(*sorted(get_cases(suite)), sep='\n')
    else:
        # start the test runner
        unittest.main(module=None, argv=["tests.main"] + other_args)