1"""Unittest main program"""
2
3import sys
4import os
5import types
6import six
7
8from unittest2 import loader, runner
9try:
10    from unittest2.signals import installHandler
11except ImportError:
12    installHandler = None
13
14__unittest = True
15
16FAILFAST = "  -f, --failfast   Stop on first failure\n"
17CATCHBREAK = "  -c, --catch      Catch control-C and display results\n"
18BUFFEROUTPUT = "  -b, --buffer     Buffer stdout and stderr during test runs\n"
19
20USAGE_AS_MAIN = """\
21Usage: %(progName)s [options] [tests]
22
23Options:
24  -h, --help       Show this message
25  -v, --verbose    Verbose output
26  -q, --quiet      Minimal output
27%(failfast)s%(catchbreak)s%(buffer)s
28Examples:
29  %(progName)s test_module                       - run tests from test_module
30  %(progName)s test_module.TestClass             - run tests from
31                                                   test_module.TestClass
32  %(progName)s test_module.TestClass.test_method - run specified test method
33
34[tests] can be a list of any number of test modules, classes and test
35methods.
36
37Alternative Usage: %(progName)s discover [options]
38
39Options:
40  -v, --verbose    Verbose output
41%(failfast)s%(catchbreak)s%(buffer)s  -s directory     Directory to start discovery ('.' default)
42  -p pattern       Pattern to match test files ('test*.py' default)
43  -t directory     Top level directory of project (default to
44                   start directory)
45
46For test discovery all test modules must be importable from the top
47level directory of the project.
48"""
49
50USAGE_FROM_MODULE = """\
51Usage: %(progName)s [options] [test] [...]
52
53Options:
54  -h, --help       Show this message
55  -v, --verbose    Verbose output
56  -q, --quiet      Minimal output
57%(failfast)s%(catchbreak)s%(buffer)s
58Examples:
59  %(progName)s                               - run default set of tests
60  %(progName)s MyTestSuite                   - run suite 'MyTestSuite'
61  %(progName)s MyTestCase.testSomething      - run MyTestCase.testSomething
62  %(progName)s MyTestCase                    - run all 'test*' test methods
63                                               in MyTestCase
64"""
65
66
67class TestProgram(object):
68    """A command-line program that runs a set of tests; this is primarily
69       for making test modules conveniently executable.
70    """
71    USAGE = USAGE_FROM_MODULE
72
73    # defaults for testing
74    failfast = catchbreak = buffer = progName = None
75
76    def __init__(self, module='__main__', defaultTest=None,
77                 argv=None, testRunner=None,
78                 testLoader=loader.defaultTestLoader, exit=True,
79                 verbosity=1, failfast=None, catchbreak=None, buffer=None):
80        if isinstance(module, six.string_types):
81            self.module = __import__(module)
82            for part in module.split('.')[1:]:
83                self.module = getattr(self.module, part)
84        else:
85            self.module = module
86        if argv is None:
87            argv = sys.argv
88
89        self.exit = exit
90        self.verbosity = verbosity
91        self.failfast = failfast
92        self.catchbreak = catchbreak
93        self.buffer = buffer
94        self.defaultTest = defaultTest
95        self.testRunner = testRunner
96        self.testLoader = testLoader
97        self.progName = os.path.basename(argv[0])
98        self.parseArgs(argv)
99        self.runTests()
100
101    def usageExit(self, msg=None):
102        if msg:
103            print(msg)
104        usage = {'progName': self.progName, 'catchbreak': '', 'failfast': '',
105                 'buffer': ''}
106        if self.failfast:
107            usage['failfast'] = FAILFAST
108        if self.catchbreak and installHandler is not None:
109            usage['catchbreak'] = CATCHBREAK
110        if self.buffer:
111            usage['buffer'] = BUFFEROUTPUT
112        print(self.USAGE % usage)
113        sys.exit(2)
114
115    def parseArgs(self, argv):
116        if len(argv) > 1 and argv[1].lower() == 'discover':
117            self._do_discovery(argv[2:])
118            return
119
120        import getopt
121        long_opts = ['help', 'verbose', 'quiet', 'failfast', 'catch', 'buffer']
122        try:
123            options, args = getopt.getopt(argv[1:], 'hHvqfcb', long_opts)
124            for opt, value in options:
125                if opt in ('-h', '-H', '--help'):
126                    self.usageExit()
127                if opt in ('-q', '--quiet'):
128                    self.verbosity = 0
129                if opt in ('-v', '--verbose'):
130                    self.verbosity = 2
131                if opt in ('-f', '--failfast'):
132                    if self.failfast is None:
133                        self.failfast = True
134                    # Should this raise an exception if -f is not valid?
135                if opt in ('-c', '--catch'):
136                    if self.catchbreak is None and installHandler is not None:
137                        self.catchbreak = True
138                    # Should this raise an exception if -c is not valid?
139                if opt in ('-b', '--buffer'):
140                    if self.buffer is None:
141                        self.buffer = True
142                    # Should this raise an exception if -b is not valid?
143            if len(args) == 0 and self.defaultTest is None:
144                # createTests will load tests from self.module
145                self.testNames = None
146            elif len(args) > 0:
147                self.testNames = args
148                if __name__ == '__main__':
149                    # to support python -m unittest ...
150                    self.module = None
151            else:
152                self.testNames = (self.defaultTest,)
153            self.createTests()
154        except getopt.error as msg:
155            self.usageExit(msg)
156
157    def createTests(self):
158        if self.testNames is None:
159            self.test = self.testLoader.loadTestsFromModule(self.module)
160        else:
161            self.test = self.testLoader.loadTestsFromNames(self.testNames,
162                                                           self.module)
163
164    def _do_discovery(self, argv, Loader=loader.TestLoader):
165        # handle command line args for test discovery
166        self.progName = '%s discover' % self.progName
167        import optparse
168        parser = optparse.OptionParser()
169        parser.prog = self.progName
170        parser.add_option('-v', '--verbose', dest='verbose', default=False,
171                          help='Verbose output', action='store_true')
172        if self.failfast:
173            parser.add_option(
174                '-f',
175                '--failfast',
176                dest='failfast',
177                default=False,
178                help='Stop on first fail or error',
179                action='store_true')
180        if self.catchbreak and installHandler is not None:
181            parser.add_option(
182                '-c',
183                '--catch',
184                dest='catchbreak',
185                default=False,
186                help='Catch ctrl-C and display results so far',
187                action='store_true')
188        if self.buffer:
189            parser.add_option('-b', '--buffer', dest='buffer', default=False,
190                              help='Buffer stdout and stderr during tests',
191                              action='store_true')
192        parser.add_option('-s', '--start-directory', dest='start', default='.',
193                          help="Directory to start discovery ('.' default)")
194        parser.add_option(
195            '-p',
196            '--pattern',
197            dest='pattern',
198            default='test*.py',
199            help="Pattern to match tests ('test*.py' default)")
200        parser.add_option(
201            '-t',
202            '--top-level-directory',
203            dest='top',
204            default=None,
205            help='Top level directory of project (defaults to start directory)')
206
207        options, args = parser.parse_args(argv)
208        if len(args) > 3:
209            self.usageExit()
210
211        for name, value in zip(('start', 'pattern', 'top'), args):
212            setattr(options, name, value)
213
214        # only set options from the parsing here
215        # if they weren't set explicitly in the constructor
216        if self.failfast is None:
217            self.failfast = options.failfast
218        if self.catchbreak is None and installHandler is not None:
219            self.catchbreak = options.catchbreak
220        if self.buffer is None:
221            self.buffer = options.buffer
222
223        if options.verbose:
224            self.verbosity = 2
225
226        start_dir = options.start
227        pattern = options.pattern
228        top_level_dir = options.top
229
230        loader = Loader()
231        self.test = loader.discover(start_dir, pattern, top_level_dir)
232
233    def runTests(self):
234        if self.catchbreak:
235            installHandler()
236        if self.testRunner is None:
237            self.testRunner = runner.TextTestRunner
238        if isinstance(self.testRunner, (type, types.ClassType)):
239            try:
240                testRunner = self.testRunner(verbosity=self.verbosity,
241                                             failfast=self.failfast,
242                                             buffer=self.buffer)
243            except TypeError:
244                # didn't accept the verbosity, buffer or failfast arguments
245                testRunner = self.testRunner()
246        else:
247            # it is assumed to be a TestRunner instance
248            testRunner = self.testRunner
249        self.result = testRunner.run(self.test)
250        if self.exit:
251            sys.exit(not self.result.wasSuccessful())
252
253main = TestProgram
254
255
256def main_():
257    TestProgram.USAGE = USAGE_AS_MAIN
258    main(module=None)
259