| #!/usr/local/bin/python2.6 -tt |
| # |
| # Copyright 2004-present Facebook. All rights reserved. |
| # |
| """ |
| This file contains the main module code for buck python test programs. |
| |
| By default, this is the main module for all python_test() rules. However, |
| rules can also specify their own custom main_module. If you write your own |
| main module, you can import this module as tools.test.stubs.fbpyunit, to access |
| any of its code to help implement your main module. |
| """ |
| |
| from __future__ import print_function |
| |
| import json |
| import optparse |
| import os |
| import sys |
| import time |
| import re |
| import unittest |
| import traceback |
| import contextlib |
| |
| |
| class TestStatus(object): |
| |
| ABORTED = 'FAILURE' |
| PASSED = 'SUCCESS' |
| FAILED = 'FAILURE' |
| EXPECTED_FAILURE = 'SUCCESS' |
| UNEXPECTED_SUCCESS = 'FAILURE' |
| SKIPPED = 'SUCCESS' |
| |
| |
| 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): |
| self._callback = callback |
| |
| def write(self, data): |
| self._callback(data) |
| |
| def flush(self): |
| pass |
| |
| def isatty(self): |
| return False |
| |
| |
| class FbJsonTestResult(unittest._TextTestResult): |
| """ |
| Our own TestResult class that outputs data in a format that can be easily |
| parsed by fbmake's test runner. |
| """ |
| |
| def __init__(self, stream, descriptions, verbosity): |
| super(FbJsonTestResult, self).__init__(stream, descriptions, verbosity) |
| self._results = [] |
| self._current_test = None |
| self._saved_stdout = sys.stdout |
| self._saved_stderr = sys.stderr |
| |
| def getResults(self): |
| return self._results |
| |
| def startTest(self, test): |
| super(FbJsonTestResult, self).startTest(test) |
| |
| sys.stdout = TeeStream(CallbackStream(self.addStdout), sys.stdout) |
| sys.stderr = TeeStream(CallbackStream(self.addStderr), 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 = '' |
| |
| def stopTest(self, test): |
| sys.stdout = self._saved_stdout |
| sys.stderr = self._saved_stderr |
| |
| super(FbJsonTestResult, self).stopTest(test) |
| |
| self._results.append({ |
| 'testCaseName': test._testMethodName, |
| 'testCase': '{0}.{1}'.format( |
| test.__class__.__module__, |
| test.__class__.__name__), |
| '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, |
| }) |
| |
| self._current_test = None |
| |
| @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(FbJsonTestResult, self).addSuccess(test) |
| self.setStatus(test, TestStatus.PASSED) |
| |
| def addError(self, test, err): |
| super(FbJsonTestResult, self).addError(test, err) |
| self.setException(test, TestStatus.ABORTED, err) |
| |
| def addFailure(self, test, err): |
| super(FbJsonTestResult, self).addFailure(test, err) |
| self.setException(test, TestStatus.FAILED, err) |
| |
| def addSkip(self, test, reason): |
| super(FbJsonTestResult, self).addSkip(test, reason) |
| self.setStatus(test, TestStatus.SKIPPED, 'Skipped: %s' % (reason,)) |
| |
| def addExpectedFailure(self, test, err): |
| super(FbJsonTestResult, self).addExpectedFailure(test, err) |
| self.setException(test, TestStatus.EXPECTED_FAILURE, err) |
| |
| def addUnexpectedSuccess(self, test): |
| super(FbJsonTestResult, self).addUnexpectedSuccess(test) |
| self.setStatus(test, TestStatus.UNEXPECTED_SUCCESS, |
| 'Unexpected success') |
| |
| def addStdout(self, val): |
| self._stdout += val |
| |
| def addStderr(self, val): |
| self._stderr += val |
| |
| |
| class FbJsonTestRunner(unittest.TextTestRunner): |
| |
| def _makeResult(self): |
| return FbJsonTestResult( |
| self.stream, |
| self.descriptions, |
| self.verbosity) |
| |
| |
| class RegexTestLoader(unittest.TestLoader): |
| |
| def __init__(self, regex=None): |
| super(unittest.TestLoader, self).__init__() |
| self.regex = regex |
| |
| def getTestCaseNames(self, testCaseClass): |
| """ |
| Return a sorted sequence of method names found within testCaseClass |
| """ |
| |
| testFnNames = super(RegexTestLoader, self).getTestCaseNames( |
| testCaseClass) |
| matched = [] |
| for attrname in testFnNames: |
| fullname = '{0}.{1}#{2}'.format( |
| testCaseClass.__class__.__module__, |
| testCaseClass.__class__.__name__, |
| attrname) |
| if self.regex is None or re.search(self.regex, fullname): |
| matched.append(attrname) |
| return matched |
| |
| |
| class Loader(object): |
| |
| def __init__(self, modules, regex=None): |
| self.modules = modules |
| self.regex = regex |
| |
| def load_all(self): |
| loader = RegexTestLoader(self.regex) |
| test_suite = unittest.TestSuite() |
| 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) |
| |
| |
| 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) |
| |
| 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( |
| '-o', '--output', |
| help='Write results to a file in a JSON format to be read by Buck') |
| op.add_option( |
| '-l', '--list-tests', action='store_true', dest='list', |
| default=False, help='List tests and exit') |
| 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( |
| '-r', '--regex', default=None, |
| help='Regex to apply to tests, to only run those tests') |
| 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 |
| |
| 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.test_args: |
| return loader.load_args(self.test_args) |
| else: |
| return loader.load_all() |
| |
| def get_test_names(self, test_suite): |
| names = [] |
| |
| for test in test_suite: |
| if isinstance(test, unittest.TestSuite): |
| names.extend(self.get_test_names(test)) |
| else: |
| names.append(str(test)) |
| |
| return names |
| |
| def run(self): |
| test_suite = self.load_tests() |
| |
| if self.options.list: |
| for test in self.get_test_names(test_suite): |
| print(test) |
| return 0 |
| 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) |
| return 0 |
| #return 0 if result.wasSuccessful() else 1 |
| |
| 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 = FbJsonTestRunner(verbosity=self.options.verbosity) |
| result = runner.run(test_suite) |
| |
| return result |
| |
| |
| def main(argv): |
| return MainProgram(sys.argv).run() |
| |
| |
| if __name__ == '__main__': |
| sys.exit(main(sys.argv)) |
