File size: 2,152 Bytes
2abfccb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
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()
|