#!/usr/bin/env python
#
# Copyright (c) Facebook, Inc. and its affiliates.
#
"""
This file contains the main module code for Python test programs.
"""
import contextlib
import ctypes
import fnmatch
import json
import logging
import optparse
import os
import platform
import re
import sys
import tempfile
import time
import traceback
import unittest
import warnings
from importlib.machinery import PathFinder
try:
from StringIO import StringIO
except ImportError:
from io import StringIO
try:
import coverage
except ImportError:
coverage = None # type: ignore
try:
from importlib.machinery import SourceFileLoader
except ImportError:
SourceFileLoader = None # type: ignore
class get_cpu_instr_counter(object):
def read(self):
# TODO
return 0
EXIT_CODE_SUCCESS = 0
EXIT_CODE_TEST_FAILURE = 70
class TestStatus(object):
ABORTED = "FAILURE"
PASSED = "SUCCESS"
FAILED = "FAILURE"
EXPECTED_FAILURE = "SUCCESS"
UNEXPECTED_SUCCESS = "FAILURE"
SKIPPED = "ASSUMPTION_VIOLATION"
class PathMatcher(object):
def __init__(self, include_patterns, omit_patterns):
self.include_patterns = include_patterns
self.omit_patterns = omit_patterns
def omit(self, path):
"""
Omit iff matches any of the omit_patterns or the include patterns are
not empty and none is matched
"""
path = os.path.realpath(path)
return any(fnmatch.fnmatch(path, p) for p in self.omit_patterns) or (
self.include_patterns
and not any(fnmatch.fnmatch(path, p) for p in self.include_patterns)
)
def include(self, path):
return not self.omit(path)
class DebugWipeFinder(PathFinder):
"""
PEP 302 finder that uses a DebugWipeLoader for all files which do not need
coverage
"""
def __init__(self, matcher):
self.matcher = matcher
def find_spec(self, fullname, path=None, target=None):
spec = super().find_spec(fullname, path=path, target=target)
if spec is None or spec.origin is None:
return None
if not spec.origin.endswith(".py"):
return None
if self.matcher.include(spec.origin):
return None
class PyVarObject(ctypes.Structure):
_fields_ = [
("ob_refcnt", ctypes.c_long),
("ob_type", ctypes.c_void_p),
("ob_size", ctypes.c_ulong),
]
class DebugWipeLoader(SourceFileLoader):
"""
PEP302 loader that zeros out debug information before execution
"""
def get_code(self, fullname):
code = super().get_code(fullname)
if code:
# Ideally we'd do
# code.co_lnotab = b''
# But code objects are READONLY. Not to worry though; we'll
# directly modify CPython's object
code_impl = PyVarObject.from_address(id(code.co_lnotab))
code_impl.ob_size = 0
return code
if isinstance(spec.loader, SourceFileLoader):
spec.loader = DebugWipeLoader(fullname, spec.origin)
return spec
def optimize_for_coverage(cov, include_patterns, omit_patterns):
"""
We get better performance if we zero out debug information for files which
we're not interested in. Only available in CPython 3.3+
"""
matcher = PathMatcher(include_patterns, omit_patterns)
if SourceFileLoader and platform.python_implementation() == "CPython":
sys.meta_path.insert(0, DebugWipeFinder(matcher))
class TeeStream(object):
def __init__(self, *streams):
self._streams = streams
def write(self, data):
for stream in self._streams:
stream.write(data)
def flush(self):
for stream in self._streams:
stream.flush()
def isatty(self):
return False
class CallbackStream(object):
def __init__(self, callback, bytes_callback=None, orig=None):
self._callback = callback
self._fileno = orig.fileno() if orig else None
# Python 3 APIs:
# - `encoding` is a string holding the encoding name
# - `errors` is a string holding the error-handling mode for encoding
# - `buffer` should look like an io.BufferedIOBase object
self.errors = orig.errors if orig else None
if bytes_callback:
# those members are only on the io.TextIOWrapper
self.encoding = orig.encoding if orig else "UTF-8"
self.buffer = CallbackStream(bytes_callback, orig=orig)
def write(self, data):
self._callback(data)
def flush(self):
pass
def isatty(self):
return False
def fileno(self):
return self._fileno
class BuckTestResult(unittest.TextTestResult):
"""
Our own TestResult class that outputs data in a format that can be easily
parsed by buck's test runner.
"""
_instr_counter = get_cpu_instr_counter()
def __init__(
self, stream, descriptions, verbosity, show_output, main_program, suite
):
super(BuckTestResult, self).__init__(stream, descriptions, verbosity)
self._main_program = main_program
self._suite = suite
self._results = []
self._current_test = None
self._saved_stdout = sys.stdout
self._saved_stderr = sys.stderr
self._show_output = show_output
def getResults(self):
return self._results
def startTest(self, test):
super(BuckTestResult, self).startTest(test)
# Pass in the real stdout and stderr filenos. We can't really do much
# here to intercept callers who directly operate on these fileno
# objects.
sys.stdout = CallbackStream(
self.addStdout, self.addStdoutBytes, orig=sys.stdout
)
sys.stderr = CallbackStream(
self.addStderr, self.addStderrBytes, orig=sys.stderr
)
self._current_test = test
self._test_start_time = time.time()
self._current_status = TestStatus.ABORTED
self._messages = []
self._stacktrace = None
self._stdout = ""
self._stderr = ""
self._start_instr_count = self._instr_counter.read()
def _find_next_test(self, suite):
"""
Find the next test that has not been run.
"""
for test in suite:
# We identify test suites by test that are iterable (as is done in
# the builtin python test harness). If we see one, recurse on it.
if hasattr(test, "__iter__"):
test = self._find_next_test(test)
# The builtin python test harness sets test references to `None`
# after they have run, so we know we've found the next test up
# if it's not `None`.
if test is not None:
return test
def stopTest(self, test):
sys.stdout = self._saved_stdout
sys.stderr = self._saved_stderr
super(BuckTestResult, self).stopTest(test)
# If a failure occurred during module/class setup, then this "test" may
# actually be a `_ErrorHolder`, which doesn't contain explicit info
# about the upcoming test. Since we really only care about the test
# name field (i.e. `_testMethodName`), we use that to detect an actual
# test cases, and fall back to looking the test up from the suite
# otherwise.
if not hasattr(test, "_testMethodName"):
test = self._find_next_test(self._suite)
result = {
"testCaseName": "{0}.{1}".format(
test.__class__.__module__, test.__class__.__name__
),
"testCase": test._testMethodName,
"type": self._current_status,
"time": int((time.time() - self._test_start_time) * 1000),
"message": os.linesep.join(self._messages),
"stacktrace": self._stacktrace,
"stdOut": self._stdout,
"stdErr": self._stderr,
}
# TestPilot supports an instruction count field.
if "TEST_PILOT" in os.environ:
result["instrCount"] = (
int(self._instr_counter.read() - self._start_instr_count),
)
self._results.append(result)
self._current_test = None
def stopTestRun(self):
cov = self._main_program.get_coverage()
if cov is not None:
self._results.append({"coverage": cov})
@contextlib.contextmanager
def _withTest(self, test):
self.startTest(test)
yield
self.stopTest(test)
def _setStatus(self, test, status, message=None, stacktrace=None):
assert test == self._current_test
self._current_status = status
self._stacktrace = stacktrace
if message is not None:
if message.endswith(os.linesep):
message = message[:-1]
self._messages.append(message)
def setStatus(self, test, status, message=None, stacktrace=None):
# addError() may be called outside of a test if one of the shared
# fixtures (setUpClass/tearDownClass/setUpModule/tearDownModule)
# throws an error.
#
# In this case, create a fake test result to record the error.
if self._current_test is None:
with self._withTest(test):
self._setStatus(test, status, message, stacktrace)
else:
self._setStatus(test, status, message, stacktrace)
def setException(self, test, status, excinfo):
exctype, value, tb = excinfo
self.setStatus(
test,
status,
"{0}: {1}".format(exctype.__name__, value),
"".join(traceback.format_tb(tb)),
)
def addSuccess(self, test):
super(BuckTestResult, self).addSuccess(test)
self.setStatus(test, TestStatus.PASSED)
def addError(self, test, err):
super(BuckTestResult, self).addError(test, err)
self.setException(test, TestStatus.ABORTED, err)
def addFailure(self, test, err):
super(BuckTestResult, self).addFailure(test, err)
self.setException(test, TestStatus.FAILED, err)
def addSkip(self, test, reason):
super(BuckTestResult, self).addSkip(test, reason)
self.setStatus(test, TestStatus.SKIPPED, "Skipped: %s" % (reason,))
def addExpectedFailure(self, test, err):
super(BuckTestResult, self).addExpectedFailure(test, err)
self.setException(test, TestStatus.EXPECTED_FAILURE, err)
def addUnexpectedSuccess(self, test):
super(BuckTestResult, self).addUnexpectedSuccess(test)
self.setStatus(test, TestStatus.UNEXPECTED_SUCCESS, "Unexpected success")
def addStdout(self, val):
self._stdout += val
if self._show_output:
self._saved_stdout.write(val)
self._saved_stdout.flush()
def addStdoutBytes(self, val):
string = val.decode("utf-8", errors="backslashreplace")
self.addStdout(string)
def addStderr(self, val):
self._stderr += val
if self._show_output:
self._saved_stderr.write(val)
self._saved_stderr.flush()
def addStderrBytes(self, val):
string = val.decode("utf-8", errors="backslashreplace")
self.addStderr(string)
class BuckTestRunner(unittest.TextTestRunner):
def __init__(self, main_program, suite, show_output=True, **kwargs):
super(BuckTestRunner, self).__init__(**kwargs)
self.show_output = show_output
self._main_program = main_program
self._suite = suite
def _makeResult(self):
return BuckTestResult(
self.stream,
self.descriptions,
self.verbosity,
self.show_output,
self._main_program,
self._suite,
)
def _format_test_name(test_class, attrname):
return "{0}.{1}.{2}".format(test_class.__module__, test_class.__name__, attrname)
class StderrLogHandler(logging.StreamHandler):
"""
This class is very similar to logging.StreamHandler, except that it
always uses the current sys.stderr object.
StreamHandler caches the current sys.stderr object when it is constructed.
This makes it behave poorly in unit tests, which may replace sys.stderr
with a StringIO buffer during tests. The StreamHandler will continue using
the old sys.stderr object instead of the desired StringIO buffer.
"""
def __init__(self):
logging.Handler.__init__(self)
@property
def stream(self):
return sys.stderr
class RegexTestLoader(unittest.TestLoader):
def __init__(self, regex=None):
self.regex = regex
super(RegexTestLoader, self).__init__()
def getTestCaseNames(self, testCaseClass):
"""
Return a sorted sequence of method names found within testCaseClass
"""
testFnNames = super(RegexTestLoader, self).getTestCaseNames(testCaseClass)
if self.regex is None:
return testFnNames
robj = re.compile(self.regex)
matched = []
for attrname in testFnNames:
fullname = _format_test_name(testCaseClass, attrname)
if robj.search(fullname):
matched.append(attrname)
return matched
class Loader(object):
suiteClass = unittest.TestSuite
def __init__(self, modules, regex=None):
self.modules = modules
self.regex = regex
def load_all(self):
loader = RegexTestLoader(self.regex)
test_suite = self.suiteClass()
for module_name in self.modules:
__import__(module_name, level=0)
module = sys.modules[module_name]
module_suite = loader.loadTestsFromModule(module)
test_suite.addTest(module_suite)
return test_suite
def load_args(self, args):
loader = RegexTestLoader(self.regex)
suites = []
for arg in args:
suite = loader.loadTestsFromName(arg)
# loadTestsFromName() can only process names that refer to
# individual test functions or modules. It can't process package
# names. If there were no module/function matches, check to see if
# this looks like a package name.
if suite.countTestCases() != 0:
suites.append(suite)
continue
# Load all modules whose name is <arg>.<something>
prefix = arg + "."
for module in self.modules:
if module.startswith(prefix):
suite = loader.loadTestsFromName(module)
suites.append(suite)
return loader.suiteClass(suites)
_COVERAGE_INI = """\
[report]
exclude_lines =
pragma: no cover
pragma: nocover
pragma:.*no${PLATFORM}
pragma:.*no${PY_IMPL}${PY_MAJOR}${PY_MINOR}
pragma:.*no${PY_IMPL}${PY_MAJOR}
pragma:.*nopy${PY_MAJOR}
pragma:.*nopy${PY_MAJOR}${PY_MINOR}
"""
class MainProgram(object):
"""
This class implements the main program. It can be subclassed by
users who wish to customize some parts of the main program.
(Adding additional command line options, customizing test loading, etc.)
"""
DEFAULT_VERBOSITY = 2
def __init__(self, argv):
self.init_option_parser()
self.parse_options(argv)
self.setup_logging()
def init_option_parser(self):
usage = "%prog [options] [TEST] ..."
op = optparse.OptionParser(usage=usage, add_help_option=False)
self.option_parser = op
op.add_option(
"--hide-output",
dest="show_output",
action="store_false",
default=True,
help="Suppress data that tests print to stdout/stderr, and only "
"show it if the test fails.",
)
op.add_option(
"-o",
"--output",
help="Write results to a file in a JSON format to be read by Buck",
)
op.add_option(
"-f",
"--failfast",
action="store_true",
default=False,
help="Stop after the first failure",
)
op.add_option(
"-l",
"--list-tests",
action="store_true",
dest="list",
default=False,
help="List tests and exit",
)
op.add_option(
"-r",
"--regex",
default=None,
help="Regex to apply to tests, to only run those tests",
)
op.add_option(
"--collect-coverage",
action="store_true",
default=False,
help="Collect test coverage information",
)
op.add_option(
"--coverage-include",
default="*",
help='File globs to include in converage (split by ",")',
)
op.add_option(
"--coverage-omit",
default="",
help='File globs to omit from converage (split by ",")',
)
op.add_option(
"--logger",
action="append",
metavar="<category>=<level>",
default=[],
help="Configure log levels for specific logger categories",
)
op.add_option(
"-q",
"--quiet",
action="count",
default=0,
help="Decrease the verbosity (may be specified multiple times)",
)
op.add_option(
"-v",
"--verbosity",
action="count",
default=self.DEFAULT_VERBOSITY,
help="Increase the verbosity (may be specified multiple times)",
)
op.add_option(
"-?", "--help", action="help", help="Show this help message and exit"
)
def parse_options(self, argv):
self.options, self.test_args = self.option_parser.parse_args(argv[1:])
self.options.verbosity -= self.options.quiet
if self.options.collect_coverage and coverage is None:
self.option_parser.error("coverage module is not available")
self.options.coverage_include = self.options.coverage_include.split(",")
if self.options.coverage_omit == "":
self.options.coverage_omit = []
else:
self.options.coverage_omit = self.options.coverage_omit.split(",")
def setup_logging(self):
# Configure the root logger to log at INFO level.
# This is similar to logging.basicConfig(), but uses our
# StderrLogHandler instead of a StreamHandler.
fmt = logging.Formatter("%(pathname)s:%(lineno)s: %(message)s")
log_handler = StderrLogHandler()
log_handler.setFormatter(fmt)
root_logger = logging.getLogger()
root_logger.addHandler(log_handler)
root_logger.setLevel(logging.INFO)
level_names = {
"debug": logging.DEBUG,
"info": logging.INFO,
"warn": logging.WARNING,
"warning": logging.WARNING,
"error": logging.ERROR,
"critical": logging.CRITICAL,
"fatal": logging.FATAL,
}
for value in self.options.logger:
parts = value.rsplit("=", 1)
if len(parts) != 2:
self.option_parser.error(
"--logger argument must be of the "
"form <name>=<level>: %s" % value
)
name = parts[0]
level_name = parts[1].lower()
level = level_names.get(level_name)
if level is None:
self.option_parser.error(
"invalid log level %r for log " "category %s" % (parts[1], name)
)
logging.getLogger(name).setLevel(level)
def create_loader(self):
import __test_modules__
return Loader(__test_modules__.TEST_MODULES, self.options.regex)
def load_tests(self):
loader = self.create_loader()
if self.options.collect_coverage:
self.start_coverage()
include = self.options.coverage_include
omit = self.options.coverage_omit
if include and "*" not in include:
optimize_for_coverage(self.cov, include, omit)
if self.test_args:
suite = loader.load_args(self.test_args)
else:
suite = loader.load_all()
if self.options.collect_coverage:
self.cov.start()
return suite
def get_tests(self, test_suite):
tests = []
for test in test_suite:
if isinstance(test, unittest.TestSuite):
tests.extend(self.get_tests(test))
else:
tests.append(test)
return tests
def run(self):
test_suite = self.load_tests()
if self.options.list:
for test in self.get_tests(test_suite):
method_name = getattr(test, "_testMethodName", "")
name = _format_test_name(test.__class__, method_name)
print(name)
return EXIT_CODE_SUCCESS
else:
result = self.run_tests(test_suite)
if self.options.output is not None:
with open(self.options.output, "w") as f:
json.dump(result.getResults(), f, indent=4, sort_keys=True)
if not result.wasSuccessful():
return EXIT_CODE_TEST_FAILURE
return EXIT_CODE_SUCCESS
def run_tests(self, test_suite):
# Install a signal handler to catch Ctrl-C and display the results
# (but only if running >2.6).
if sys.version_info[0] > 2 or sys.version_info[1] > 6:
unittest.installHandler()
# Run the tests
runner = BuckTestRunner(
self,
test_suite,
verbosity=self.options.verbosity,
show_output=self.options.show_output,
)
result = runner.run(test_suite)
if self.options.collect_coverage and self.options.show_output:
self.cov.stop()
try:
self.cov.report(file=sys.stdout)
except coverage.misc.CoverageException:
print("No lines were covered, potentially restricted by file filters")
return result
def get_abbr_impl(self):
"""Return abbreviated implementation name."""
impl = platform.python_implementation()
if impl == "PyPy":
return "pp"
elif impl == "Jython":
return "jy"
elif impl == "IronPython":
return "ip"
elif impl == "CPython":
return "cp"
else:
raise RuntimeError("unknown python runtime")
def start_coverage(self):
if not self.options.collect_coverage:
return
with tempfile.NamedTemporaryFile("w", delete=False) as coverage_ini:
coverage_ini.write(_COVERAGE_INI)
self._coverage_ini_path = coverage_ini.name
# Keep the original working dir in case tests use os.chdir
self._original_working_dir = os.getcwd()
# for coverage config ignores by platform/python version
os.environ["PLATFORM"] = sys.platform
os.environ["PY_IMPL"] = self.get_abbr_impl()
os.environ["PY_MAJOR"] = str(sys.version_info.major)
os.environ["PY_MINOR"] = str(sys.version_info.minor)
self.cov = coverage.Coverage(
include=self.options.coverage_include,
omit=self.options.coverage_omit,
config_file=coverage_ini.name,
)
self.cov.erase()
self.cov.start()
def get_coverage(self):
if not self.options.collect_coverage:
return None
try:
os.remove(self._coverage_ini_path)
except OSError:
pass # Better to litter than to fail the test
# Switch back to the original working directory.
os.chdir(self._original_working_dir)
result = {}
self.cov.stop()
try:
f = StringIO()
self.cov.report(file=f)
lines = f.getvalue().split("\n")
except coverage.misc.CoverageException:
# Nothing was covered. That's fine by us
return result
# N.B.: the format of the coverage library's output differs
# depending on whether one or more files are in the results
for line in lines[2:]:
if line.strip("-") == "":
break
r = line.split()[0]
analysis = self.cov.analysis2(r)
covString = self.convert_to_diff_cov_str(analysis)
if covString:
result[r] = covString
return result
def convert_to_diff_cov_str(self, analysis):
# Info on the format of analysis:
# http://nedbatchelder.com/code/coverage/api.html
if not analysis:
return None
numLines = max(
analysis[1][-1] if len(analysis[1]) else 0,
analysis[2][-1] if len(analysis[2]) else 0,
analysis[3][-1] if len(analysis[3]) else 0,
)
lines = ["N"] * numLines
for l in analysis[1]:
lines[l - 1] = "C"
for l in analysis[2]:
lines[l - 1] = "X"
for l in analysis[3]:
lines[l - 1] = "U"
return "".join(lines)
def main(argv):
return MainProgram(sys.argv).run()
if __name__ == "__main__":
sys.exit(main(sys.argv))