af9dcada53ff9b24d564034436d16245e6a16ff3
[third_party/subunit] / python / subunit / _output.py
1 #  subunit: extensions to python unittest to get test results from subprocesses.
2 #  Copyright (C) 2013 Subunit Contributors
3 #
4 #  Licensed under either the Apache License, Version 2.0 or the BSD 3-clause
5 #  license at the users choice. A copy of both licenses are available in the
6 #  project source as Apache-2.0 and BSD. You may not use this file except in
7 #  compliance with one of these two licences.
8 #
9 #  Unless required by applicable law or agreed to in writing, software
10 #  distributed under these licenses is distributed on an "AS IS" BASIS, WITHOUT
11 #  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
12 #  license you chose for the specific language governing permissions and
13 #  limitations under that license.
14 #
15
16 import datetime
17 from functools import partial
18 from optparse import (
19     OptionGroup,
20     OptionParser,
21     OptionValueError,
22 )
23 import sys
24
25 from subunit.iso8601 import UTC
26 from subunit.v2 import StreamResultToBytes
27
28
29 _FINAL_ACTIONS = frozenset([
30     'exists',
31     'fail',
32     'skip',
33     'success',
34     'uxsuccess',
35     'xfail',
36 ])
37 _ALL_ACTIONS = _FINAL_ACTIONS.union(['inprogress'])
38 _CHUNK_SIZE=3670016 # 3.5 MiB
39
40
41 def output_main():
42     args = parse_arguments()
43     output = StreamResultToBytes(sys.stdout)
44     generate_stream_results(args, output)
45     return 0
46
47
48 def parse_arguments(args=None, ParserClass=OptionParser):
49     """Parse arguments from the command line.
50
51     If specified, args must be a list of strings, similar to sys.argv[1:].
52
53     ParserClass may be specified to override the class we use to parse the
54     command-line arguments. This is useful for testing.
55     """
56     parser = ParserClass(
57         prog="subunit-output",
58         description="A tool to generate a subunit v2 result byte-stream",
59         usage="subunit-output [-h] [status TEST_ID] [options]",
60     )
61     parser.set_default('tags', None)
62     parser.set_default('test_id', None)
63
64     status_commands = OptionGroup(
65         parser,
66         "Status Commands",
67         "These options report the status of a test. TEST_ID must be a string "
68             "that uniquely identifies the test."
69     )
70     for action_name in _ALL_ACTIONS:
71         status_commands.add_option(
72             "--%s" % action_name,
73             nargs=1,
74             action="callback",
75             callback=set_status_cb,
76             callback_args=(action_name,),
77             dest="action",
78             metavar="TEST_ID",
79             help="Report a test status."
80         )
81     parser.add_option_group(status_commands)
82
83     file_commands = OptionGroup(
84         parser,
85         "File Options",
86         "These options control attaching data to a result stream. They can "
87             "either be specified with a status command, in which case the file "
88             "is attached to the test status, or by themselves, in which case "
89             "the file is attached to the stream (and not associated with any "
90             "test id)."
91     )
92     file_commands.add_option(
93         "--attach-file",
94         help="Attach a file to the result stream for this test. If '-' is "
95             "specified, stdin will be read instead. In this case, the file "
96             "name will be set to 'stdin' (but can still be overridden with "
97             "the --file-name option)."
98     )
99     file_commands.add_option(
100         "--file-name",
101         help="The name to give this file attachment. If not specified, the "
102             "name of the file on disk will be used, or 'stdin' in the case "
103             "where '-' was passed to the '--attach-file' argument. This option"
104             " may only be specified when '--attach-file' is specified.",
105         )
106     file_commands.add_option(
107         "--mimetype",
108         help="The mime type to send with this file. This is only used if the "
109             "--attach-file argument is used. This argument is optional. If it "
110             "is not specified, the file will be sent wihtout a mime type. This "
111             "option may only be specified when '--attach-file' is specified.",
112         default=None
113     )
114     parser.add_option_group(file_commands)
115
116     parser.add_option(
117         "--tags",
118         help="A comma-separated list of tags to associate with a test. This "
119             "option may only be used with a status command.",
120         action="callback",
121         callback=set_tags_cb,
122         default=[]
123     )
124
125     (options, args) = parser.parse_args(args)
126     if options.mimetype and not options.attach_file:
127         parser.error("Cannot specify --mimetype without --attach-file")
128     if options.file_name and not options.attach_file:
129         parser.error("Cannot specify --file-name without --attach-file")
130     if options.attach_file:
131         if options.attach_file == '-':
132             if not options.file_name:
133                 options.file_name = 'stdin'
134             if sys.version[0] >= '3':
135                 options.attach_file = getattr(sys.stdin, 'buffer', sys.stdin)
136             else:
137                 options.attach_file = sys.stdin
138         else:
139             try:
140                 options.attach_file = open(options.attach_file, 'rb')
141             except IOError as e:
142                 parser.error("Cannot open %s (%s)" % (options.attach_file, e.strerror))
143     if options.tags and not options.action:
144         parser.error("Cannot specify --tags without a status command")
145     if not (options.attach_file or options.action):
146         parser.error("Must specify either --attach-file or a status command")
147
148     return options
149
150
151 def set_status_cb(option, opt_str, value, parser, status_name):
152     if getattr(parser.values, "action", None) is not None:
153         raise OptionValueError("argument %s: Only one status may be specified at once." % opt_str)
154
155     if len(parser.rargs) == 0:
156         raise OptionValueError("argument %s: must specify a single TEST_ID." % opt_str)
157     parser.values.action = status_name
158     parser.values.test_id = parser.rargs.pop(0)
159
160
161 def set_tags_cb(option, opt_str, value, parser):
162     if not parser.rargs:
163         raise OptionValueError("Must specify at least one tag with --tags")
164     parser.values.tags = parser.rargs.pop(0).split(',')
165
166
167 def generate_stream_results(args, output_writer):
168     output_writer.startTestRun()
169
170     if args.attach_file:
171         reader = partial(args.attach_file.read, _CHUNK_SIZE)
172         this_file_hunk = reader()
173         next_file_hunk = reader()
174
175     is_first_packet = True
176     is_last_packet = False
177     while not is_last_packet:
178         write_status = output_writer.status
179
180         if is_first_packet:
181             if args.attach_file:
182                 if args.mimetype:
183                     write_status = partial(write_status, mime_type=args.mimetype)
184             if args.tags:
185                 write_status = partial(write_status, test_tags=args.tags)
186             write_status = partial(write_status, timestamp=create_timestamp())
187             if args.action not in _FINAL_ACTIONS:
188                 write_status = partial(write_status, test_status=args.action)
189             is_first_packet = False
190
191         if args.attach_file:
192             filename = args.file_name or args.attach_file.name
193             write_status = partial(write_status, file_name=filename, file_bytes=this_file_hunk)
194             if next_file_hunk == b'':
195                 write_status = partial(write_status, eof=True)
196                 is_last_packet = True
197             else:
198                 this_file_hunk = next_file_hunk
199                 next_file_hunk = reader()
200         else:
201             is_last_packet = True
202
203         if args.test_id:
204             write_status = partial(write_status, test_id=args.test_id)
205
206         if is_last_packet:
207             write_status = partial(write_status, eof=True)
208             if args.action in _FINAL_ACTIONS:
209                 write_status = partial(write_status, test_status=args.action)
210
211         write_status()
212
213     output_writer.stopTestRun()
214
215
216 def create_timestamp():
217     return datetime.datetime.now(UTC)