tests/common_util.py

import sys
import argparse
import os
import unittest
import subprocess

test_dir = os.path.dirname(os.path.realpath(__file__))

parser = argparse.ArgumentParser(add_help=False)
parser.add_argument('-s', '--save-result', nargs='?', type=str, default=None)
args, remaining = parser.parse_known_args()

UNITTEST_ARGS = [sys.argv[0]] + remaining


def wait_for_process(p):
    try:
        return p.wait()
    except KeyboardInterrupt:
        # Give `p` a chance to handle KeyboardInterrupt. Without this,
        # `pytest` can't print errors it collected so far upon KeyboardInterrupt.
        exit_status = p.wait(timeout=5)
        if exit_status is not None:
            return exit_status
        else:
            p.kill()
            raise
    except:  # noqa E722, copied from python core library
        p.kill()
        raise
    finally:
        # Always call p.wait() to ensure exit
        p.wait()


def shell(command, cwd=None, env=None):
    sys.stdout.flush()
    sys.stderr.flush()
    # The following cool snippet is copied from Py3 core library subprocess.call
    # only the with
    #   1. `except KeyboardInterrupt` block added for SIGINT handling.
    #   2. In Py2, subprocess.Popen doesn't return a context manager, so we do
    #      `p.wait()` in a `final` block for the code to be portable.
    #
    # https://github.com/python/cpython/blob/71b6c1af727fbe13525fb734568057d78cea33f3/Lib/subprocess.py#L309-L323
    # assert not isinstance(command, torch._six.string_classes), "Command to shell should be a list or tuple of tokens"
    p = subprocess.Popen(command, universal_newlines=True, cwd=cwd, env=env)
    return wait_for_process(p)


def run_test(argv=UNITTEST_ARGS):
    if args.save_result is not None:
        test_report_path = test_dir + "/" + args.save_result
        with open(test_report_path, "a") as report_file:
            runner = unittest.TextTestRunner(stream=report_file, verbosity=2)
            unittest.main(argv=argv, testRunner=runner)
    else:
        runner = unittest.TextTestRunner(verbosity=2)
        unittest.main(argv=argv, testRunner=runner)


if __name__ == "__main__":
    run_test()