1"""
2A nose plugin that saves test stdout into files.  Useful for situations where
3lots of tests are breaking and you don't want to have to scroll through a
4bunch of output to find individual test output.
5"""
6
7import sys
8stderr = sys.stderr
9
10import nose.case
11
12import logging
13import os
14from nose.plugins.base import Plugin
15import string
16
17try:
18    from StringIO import StringIO as p_StringO
19except:
20    from io import StringIO as p_StringO
21
22import traceback
23
24def write_test_output(test, output, dirname, prefix=''):
25    testname = calc_testname(test)
26
27    filename = '%s%s.txt' % (prefix, testname,)
28    if dirname:
29        filename = os.path.join(dirname, filename)
30
31    fp = open(filename, 'w')
32    fp.write(output)
33    fp.close()
34
35log = logging.getLogger(__name__)
36
37# Based on code from http://stackoverflow.com/a/295146/28275
38valid_chars = frozenset("-_.()%s%s" % (string.ascii_letters, string.digits))
39def sanitize_filename(name):
40    return ''.join(c for c in name if c in valid_chars)
41
42def calc_testname(test):
43    # For errors at module-level nose passes a context instead of a test case.
44    if isinstance(test, nose.suite.ContextSuite):
45        name = test.context.__name__
46    else:
47        name = str(test)
48    if ' ' in name:
49        name = name.split(' ')[1]
50
51    return sanitize_filename(name)
52
53def get_stdout():
54    if isinstance(sys.stdout, p_StringO):
55        return sys.stdout.getvalue()
56    return None
57
58class OutputSave(Plugin):
59    def __init__(self):
60        Plugin.__init__(self)
61        self.testname = None
62
63    def add_options(self, parser, env=os.environ):
64        env_opt = 'NOSE_WITH_%s' % self.name.upper()
65        env_opt.replace('-', '_')
66        parser.add_option("--with-%s" % self.name,
67                          action="store_true",
68                          dest=self.enableOpt,
69                          default=env.get(env_opt),
70                          help="Enable plugin %s: %s [%s]" %
71                          (self.__class__.__name__, self.help(), env_opt))
72
73        parser.add_option('--omit-success-output',
74                          action="store_true",
75                          dest="omit_success_output",
76                          help = 'do not save output from successful tests',
77                          default=False)
78
79        parser.add_option('--save-directory',
80                          action="store",
81                          dest="save_directory",
82                          help="save output files to this directory",
83                          default='')
84
85    def configure(self, options, config):
86        self.conf = config
87        config.capture = False
88
89        if hasattr(options, self.enableOpt):
90            self.enabled = getattr(options, self.enableOpt)
91
92        if hasattr(options, 'omit_success_output'):
93            self.omit_success = getattr(options, 'omit_success_output')
94
95        if hasattr(options, 'save_directory',):
96            self.save_directory = getattr(options, 'save_directory')
97            self.save_directory = os.path.abspath(self.save_directory)
98
99            try:
100                os.makedirs(self.save_directory)
101            except OSError:
102                pass
103
104    def addError(self, test, err):
105        exception_text = traceback.format_exception(*err)
106        exception_text = "".join(exception_text)
107
108        capt = get_stdout()
109        if capt is None:
110            capt = exception_text
111        else:
112            capt += exception_text
113
114        write_test_output(test, capt, self.save_directory, prefix='error-')
115
116    def addFailure(self, test, err):
117        exception_text = traceback.format_exception(*err)
118        exception_text = "".join(exception_text)
119
120        capt = get_stdout()
121        if capt is None:
122            capt = exception_text
123        else:
124            capt += exception_text
125
126        write_test_output(test, capt, self.save_directory,
127                          prefix='failure-')
128
129    def addSuccess(self, test):
130        capt = get_stdout()
131        if capt is None:
132            capt = ""
133
134        if not self.omit_success:
135            write_test_output(test, capt, self.save_directory,
136                              prefix='success-')
137