iotests: Use // for Python integer division
[qemu/ar7.git] / tests / qemu-iotests / iotests.py
blob7ca94e927898b8fa198daa3a4508eb13827a7585
1 from __future__ import print_function
2 # Common utilities and Python wrappers for qemu-iotests
4 # Copyright (C) 2012 IBM Corp.
6 # This program is free software; you can redistribute it and/or modify
7 # it under the terms of the GNU General Public License as published by
8 # the Free Software Foundation; either version 2 of the License, or
9 # (at your option) any later version.
11 # This program is distributed in the hope that it will be useful,
12 # but WITHOUT ANY WARRANTY; without even the implied warranty of
13 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14 # GNU General Public License for more details.
16 # You should have received a copy of the GNU General Public License
17 # along with this program. If not, see <http://www.gnu.org/licenses/>.
20 import errno
21 import os
22 import re
23 import subprocess
24 import string
25 import unittest
26 import sys
27 import struct
28 import json
29 import signal
30 import logging
31 import atexit
33 sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', 'scripts'))
34 import qtest
37 # This will not work if arguments contain spaces but is necessary if we
38 # want to support the override options that ./check supports.
39 qemu_img_args = [os.environ.get('QEMU_IMG_PROG', 'qemu-img')]
40 if os.environ.get('QEMU_IMG_OPTIONS'):
41 qemu_img_args += os.environ['QEMU_IMG_OPTIONS'].strip().split(' ')
43 qemu_io_args = [os.environ.get('QEMU_IO_PROG', 'qemu-io')]
44 if os.environ.get('QEMU_IO_OPTIONS'):
45 qemu_io_args += os.environ['QEMU_IO_OPTIONS'].strip().split(' ')
47 qemu_nbd_args = [os.environ.get('QEMU_NBD_PROG', 'qemu-nbd')]
48 if os.environ.get('QEMU_NBD_OPTIONS'):
49 qemu_nbd_args += os.environ['QEMU_NBD_OPTIONS'].strip().split(' ')
51 qemu_prog = os.environ.get('QEMU_PROG', 'qemu')
52 qemu_opts = os.environ.get('QEMU_OPTIONS', '').strip().split(' ')
54 imgfmt = os.environ.get('IMGFMT', 'raw')
55 imgproto = os.environ.get('IMGPROTO', 'file')
56 test_dir = os.environ.get('TEST_DIR')
57 output_dir = os.environ.get('OUTPUT_DIR', '.')
58 cachemode = os.environ.get('CACHEMODE')
59 qemu_default_machine = os.environ.get('QEMU_DEFAULT_MACHINE')
61 socket_scm_helper = os.environ.get('SOCKET_SCM_HELPER', 'socket_scm_helper')
62 debug = False
64 luks_default_secret_object = 'secret,id=keysec0,data=' + \
65 os.environ['IMGKEYSECRET']
66 luks_default_key_secret_opt = 'key-secret=keysec0'
69 def qemu_img(*args):
70 '''Run qemu-img and return the exit code'''
71 devnull = open('/dev/null', 'r+')
72 exitcode = subprocess.call(qemu_img_args + list(args), stdin=devnull, stdout=devnull)
73 if exitcode < 0:
74 sys.stderr.write('qemu-img received signal %i: %s\n' % (-exitcode, ' '.join(qemu_img_args + list(args))))
75 return exitcode
77 def qemu_img_create(*args):
78 args = list(args)
80 # default luks support
81 if '-f' in args and args[args.index('-f') + 1] == 'luks':
82 if '-o' in args:
83 i = args.index('-o')
84 if 'key-secret' not in args[i + 1]:
85 args[i + 1].append(luks_default_key_secret_opt)
86 args.insert(i + 2, '--object')
87 args.insert(i + 3, luks_default_secret_object)
88 else:
89 args = ['-o', luks_default_key_secret_opt,
90 '--object', luks_default_secret_object] + args
92 args.insert(0, 'create')
94 return qemu_img(*args)
96 def qemu_img_verbose(*args):
97 '''Run qemu-img without suppressing its output and return the exit code'''
98 exitcode = subprocess.call(qemu_img_args + list(args))
99 if exitcode < 0:
100 sys.stderr.write('qemu-img received signal %i: %s\n' % (-exitcode, ' '.join(qemu_img_args + list(args))))
101 return exitcode
103 def qemu_img_pipe(*args):
104 '''Run qemu-img and return its output'''
105 subp = subprocess.Popen(qemu_img_args + list(args),
106 stdout=subprocess.PIPE,
107 stderr=subprocess.STDOUT,
108 universal_newlines=True)
109 exitcode = subp.wait()
110 if exitcode < 0:
111 sys.stderr.write('qemu-img received signal %i: %s\n' % (-exitcode, ' '.join(qemu_img_args + list(args))))
112 return subp.communicate()[0]
114 def img_info_log(filename, filter_path=None, imgopts=False, extra_args=[]):
115 args = [ 'info' ]
116 if imgopts:
117 args.append('--image-opts')
118 else:
119 args += [ '-f', imgfmt ]
120 args += extra_args
121 args.append(filename)
123 output = qemu_img_pipe(*args)
124 if not filter_path:
125 filter_path = filename
126 log(filter_img_info(output, filter_path))
128 def qemu_io(*args):
129 '''Run qemu-io and return the stdout data'''
130 args = qemu_io_args + list(args)
131 subp = subprocess.Popen(args, stdout=subprocess.PIPE,
132 stderr=subprocess.STDOUT,
133 universal_newlines=True)
134 exitcode = subp.wait()
135 if exitcode < 0:
136 sys.stderr.write('qemu-io received signal %i: %s\n' % (-exitcode, ' '.join(args)))
137 return subp.communicate()[0]
139 def qemu_io_silent(*args):
140 '''Run qemu-io and return the exit code, suppressing stdout'''
141 args = qemu_io_args + list(args)
142 exitcode = subprocess.call(args, stdout=open('/dev/null', 'w'))
143 if exitcode < 0:
144 sys.stderr.write('qemu-io received signal %i: %s\n' %
145 (-exitcode, ' '.join(args)))
146 return exitcode
149 class QemuIoInteractive:
150 def __init__(self, *args):
151 self.args = qemu_io_args + list(args)
152 self._p = subprocess.Popen(self.args, stdin=subprocess.PIPE,
153 stdout=subprocess.PIPE,
154 stderr=subprocess.STDOUT,
155 universal_newlines=True)
156 assert self._p.stdout.read(9) == 'qemu-io> '
158 def close(self):
159 self._p.communicate('q\n')
161 def _read_output(self):
162 pattern = 'qemu-io> '
163 n = len(pattern)
164 pos = 0
165 s = []
166 while pos != n:
167 c = self._p.stdout.read(1)
168 # check unexpected EOF
169 assert c != ''
170 s.append(c)
171 if c == pattern[pos]:
172 pos += 1
173 else:
174 pos = 0
176 return ''.join(s[:-n])
178 def cmd(self, cmd):
179 # quit command is in close(), '\n' is added automatically
180 assert '\n' not in cmd
181 cmd = cmd.strip()
182 assert cmd != 'q' and cmd != 'quit'
183 self._p.stdin.write(cmd + '\n')
184 self._p.stdin.flush()
185 return self._read_output()
188 def qemu_nbd(*args):
189 '''Run qemu-nbd in daemon mode and return the parent's exit code'''
190 return subprocess.call(qemu_nbd_args + ['--fork'] + list(args))
192 def compare_images(img1, img2, fmt1=imgfmt, fmt2=imgfmt):
193 '''Return True if two image files are identical'''
194 return qemu_img('compare', '-f', fmt1,
195 '-F', fmt2, img1, img2) == 0
197 def create_image(name, size):
198 '''Create a fully-allocated raw image with sector markers'''
199 file = open(name, 'wb')
200 i = 0
201 while i < size:
202 sector = struct.pack('>l504xl', i // 512, i // 512)
203 file.write(sector)
204 i = i + 512
205 file.close()
207 def image_size(img):
208 '''Return image's virtual size'''
209 r = qemu_img_pipe('info', '--output=json', '-f', imgfmt, img)
210 return json.loads(r)['virtual-size']
212 test_dir_re = re.compile(r"%s" % test_dir)
213 def filter_test_dir(msg):
214 return test_dir_re.sub("TEST_DIR", msg)
216 win32_re = re.compile(r"\r")
217 def filter_win32(msg):
218 return win32_re.sub("", msg)
220 qemu_io_re = re.compile(r"[0-9]* ops; [0-9\/:. sec]* \([0-9\/.inf]* [EPTGMKiBbytes]*\/sec and [0-9\/.inf]* ops\/sec\)")
221 def filter_qemu_io(msg):
222 msg = filter_win32(msg)
223 return qemu_io_re.sub("X ops; XX:XX:XX.X (XXX YYY/sec and XXX ops/sec)", msg)
225 chown_re = re.compile(r"chown [0-9]+:[0-9]+")
226 def filter_chown(msg):
227 return chown_re.sub("chown UID:GID", msg)
229 def filter_qmp_event(event):
230 '''Filter a QMP event dict'''
231 event = dict(event)
232 if 'timestamp' in event:
233 event['timestamp']['seconds'] = 'SECS'
234 event['timestamp']['microseconds'] = 'USECS'
235 return event
237 def filter_testfiles(msg):
238 prefix = os.path.join(test_dir, "%s-" % (os.getpid()))
239 return msg.replace(prefix, 'TEST_DIR/PID-')
241 def filter_img_info(output, filename):
242 lines = []
243 for line in output.split('\n'):
244 if 'disk size' in line or 'actual-size' in line:
245 continue
246 line = line.replace(filename, 'TEST_IMG') \
247 .replace(imgfmt, 'IMGFMT')
248 line = re.sub('iters: [0-9]+', 'iters: XXX', line)
249 line = re.sub('uuid: [-a-f0-9]+', 'uuid: XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX', line)
250 lines.append(line)
251 return '\n'.join(lines)
253 def log(msg, filters=[]):
254 for flt in filters:
255 msg = flt(msg)
256 print(msg)
258 class Timeout:
259 def __init__(self, seconds, errmsg = "Timeout"):
260 self.seconds = seconds
261 self.errmsg = errmsg
262 def __enter__(self):
263 signal.signal(signal.SIGALRM, self.timeout)
264 signal.setitimer(signal.ITIMER_REAL, self.seconds)
265 return self
266 def __exit__(self, type, value, traceback):
267 signal.setitimer(signal.ITIMER_REAL, 0)
268 return False
269 def timeout(self, signum, frame):
270 raise Exception(self.errmsg)
273 class FilePath(object):
274 '''An auto-generated filename that cleans itself up.
276 Use this context manager to generate filenames and ensure that the file
277 gets deleted::
279 with TestFilePath('test.img') as img_path:
280 qemu_img('create', img_path, '1G')
281 # migration_sock_path is automatically deleted
283 def __init__(self, name):
284 filename = '{0}-{1}'.format(os.getpid(), name)
285 self.path = os.path.join(test_dir, filename)
287 def __enter__(self):
288 return self.path
290 def __exit__(self, exc_type, exc_val, exc_tb):
291 try:
292 os.remove(self.path)
293 except OSError:
294 pass
295 return False
298 def file_path_remover():
299 for path in reversed(file_path_remover.paths):
300 try:
301 os.remove(path)
302 except OSError:
303 pass
306 def file_path(*names):
307 ''' Another way to get auto-generated filename that cleans itself up.
309 Use is as simple as:
311 img_a, img_b = file_path('a.img', 'b.img')
312 sock = file_path('socket')
315 if not hasattr(file_path_remover, 'paths'):
316 file_path_remover.paths = []
317 atexit.register(file_path_remover)
319 paths = []
320 for name in names:
321 filename = '{0}-{1}'.format(os.getpid(), name)
322 path = os.path.join(test_dir, filename)
323 file_path_remover.paths.append(path)
324 paths.append(path)
326 return paths[0] if len(paths) == 1 else paths
328 def remote_filename(path):
329 if imgproto == 'file':
330 return path
331 elif imgproto == 'ssh':
332 return "ssh://127.0.0.1%s" % (path)
333 else:
334 raise Exception("Protocol %s not supported" % (imgproto))
336 class VM(qtest.QEMUQtestMachine):
337 '''A QEMU VM'''
339 def __init__(self, path_suffix=''):
340 name = "qemu%s-%d" % (path_suffix, os.getpid())
341 super(VM, self).__init__(qemu_prog, qemu_opts, name=name,
342 test_dir=test_dir,
343 socket_scm_helper=socket_scm_helper)
344 self._num_drives = 0
346 def add_object(self, opts):
347 self._args.append('-object')
348 self._args.append(opts)
349 return self
351 def add_device(self, opts):
352 self._args.append('-device')
353 self._args.append(opts)
354 return self
356 def add_drive_raw(self, opts):
357 self._args.append('-drive')
358 self._args.append(opts)
359 return self
361 def add_drive(self, path, opts='', interface='virtio', format=imgfmt):
362 '''Add a virtio-blk drive to the VM'''
363 options = ['if=%s' % interface,
364 'id=drive%d' % self._num_drives]
366 if path is not None:
367 options.append('file=%s' % path)
368 options.append('format=%s' % format)
369 options.append('cache=%s' % cachemode)
371 if opts:
372 options.append(opts)
374 if format == 'luks' and 'key-secret' not in opts:
375 # default luks support
376 if luks_default_secret_object not in self._args:
377 self.add_object(luks_default_secret_object)
379 options.append(luks_default_key_secret_opt)
381 self._args.append('-drive')
382 self._args.append(','.join(options))
383 self._num_drives += 1
384 return self
386 def add_blockdev(self, opts):
387 self._args.append('-blockdev')
388 if isinstance(opts, str):
389 self._args.append(opts)
390 else:
391 self._args.append(','.join(opts))
392 return self
394 def add_incoming(self, addr):
395 self._args.append('-incoming')
396 self._args.append(addr)
397 return self
399 def pause_drive(self, drive, event=None):
400 '''Pause drive r/w operations'''
401 if not event:
402 self.pause_drive(drive, "read_aio")
403 self.pause_drive(drive, "write_aio")
404 return
405 self.qmp('human-monitor-command',
406 command_line='qemu-io %s "break %s bp_%s"' % (drive, event, drive))
408 def resume_drive(self, drive):
409 self.qmp('human-monitor-command',
410 command_line='qemu-io %s "remove_break bp_%s"' % (drive, drive))
412 def hmp_qemu_io(self, drive, cmd):
413 '''Write to a given drive using an HMP command'''
414 return self.qmp('human-monitor-command',
415 command_line='qemu-io %s "%s"' % (drive, cmd))
417 def flatten_qmp_object(self, obj, output=None, basestr=''):
418 if output is None:
419 output = dict()
420 if isinstance(obj, list):
421 for i in range(len(obj)):
422 self.flatten_qmp_object(obj[i], output, basestr + str(i) + '.')
423 elif isinstance(obj, dict):
424 for key in obj:
425 self.flatten_qmp_object(obj[key], output, basestr + key + '.')
426 else:
427 output[basestr[:-1]] = obj # Strip trailing '.'
428 return output
430 def qmp_to_opts(self, obj):
431 obj = self.flatten_qmp_object(obj)
432 output_list = list()
433 for key in obj:
434 output_list += [key + '=' + obj[key]]
435 return ','.join(output_list)
437 def get_qmp_events_filtered(self, wait=True):
438 result = []
439 for ev in self.get_qmp_events(wait=wait):
440 result.append(filter_qmp_event(ev))
441 return result
443 def qmp_log(self, cmd, filters=[filter_testfiles], **kwargs):
444 logmsg = "{'execute': '%s', 'arguments': %s}" % (cmd, kwargs)
445 log(logmsg, filters)
446 result = self.qmp(cmd, **kwargs)
447 log(str(result), filters)
448 return result
450 def run_job(self, job, auto_finalize=True, auto_dismiss=False):
451 while True:
452 for ev in self.get_qmp_events_filtered(wait=True):
453 if ev['event'] == 'JOB_STATUS_CHANGE':
454 status = ev['data']['status']
455 if status == 'aborting':
456 result = self.qmp('query-jobs')
457 for j in result['return']:
458 if j['id'] == job:
459 log('Job failed: %s' % (j['error']))
460 elif status == 'pending' and not auto_finalize:
461 self.qmp_log('job-finalize', id=job)
462 elif status == 'concluded' and not auto_dismiss:
463 self.qmp_log('job-dismiss', id=job)
464 elif status == 'null':
465 return
466 else:
467 iotests.log(ev)
470 index_re = re.compile(r'([^\[]+)\[([^\]]+)\]')
472 class QMPTestCase(unittest.TestCase):
473 '''Abstract base class for QMP test cases'''
475 def dictpath(self, d, path):
476 '''Traverse a path in a nested dict'''
477 for component in path.split('/'):
478 m = index_re.match(component)
479 if m:
480 component, idx = m.groups()
481 idx = int(idx)
483 if not isinstance(d, dict) or component not in d:
484 self.fail('failed path traversal for "%s" in "%s"' % (path, str(d)))
485 d = d[component]
487 if m:
488 if not isinstance(d, list):
489 self.fail('path component "%s" in "%s" is not a list in "%s"' % (component, path, str(d)))
490 try:
491 d = d[idx]
492 except IndexError:
493 self.fail('invalid index "%s" in path "%s" in "%s"' % (idx, path, str(d)))
494 return d
496 def assert_qmp_absent(self, d, path):
497 try:
498 result = self.dictpath(d, path)
499 except AssertionError:
500 return
501 self.fail('path "%s" has value "%s"' % (path, str(result)))
503 def assert_qmp(self, d, path, value):
504 '''Assert that the value for a specific path in a QMP dict matches'''
505 result = self.dictpath(d, path)
506 self.assertEqual(result, value, 'values not equal "%s" and "%s"' % (str(result), str(value)))
508 def assert_no_active_block_jobs(self):
509 result = self.vm.qmp('query-block-jobs')
510 self.assert_qmp(result, 'return', [])
512 def assert_has_block_node(self, node_name=None, file_name=None):
513 """Issue a query-named-block-nodes and assert node_name and/or
514 file_name is present in the result"""
515 def check_equal_or_none(a, b):
516 return a == None or b == None or a == b
517 assert node_name or file_name
518 result = self.vm.qmp('query-named-block-nodes')
519 for x in result["return"]:
520 if check_equal_or_none(x.get("node-name"), node_name) and \
521 check_equal_or_none(x.get("file"), file_name):
522 return
523 self.assertTrue(False, "Cannot find %s %s in result:\n%s" % \
524 (node_name, file_name, result))
526 def assert_json_filename_equal(self, json_filename, reference):
527 '''Asserts that the given filename is a json: filename and that its
528 content is equal to the given reference object'''
529 self.assertEqual(json_filename[:5], 'json:')
530 self.assertEqual(self.vm.flatten_qmp_object(json.loads(json_filename[5:])),
531 self.vm.flatten_qmp_object(reference))
533 def cancel_and_wait(self, drive='drive0', force=False, resume=False):
534 '''Cancel a block job and wait for it to finish, returning the event'''
535 result = self.vm.qmp('block-job-cancel', device=drive, force=force)
536 self.assert_qmp(result, 'return', {})
538 if resume:
539 self.vm.resume_drive(drive)
541 cancelled = False
542 result = None
543 while not cancelled:
544 for event in self.vm.get_qmp_events(wait=True):
545 if event['event'] == 'BLOCK_JOB_COMPLETED' or \
546 event['event'] == 'BLOCK_JOB_CANCELLED':
547 self.assert_qmp(event, 'data/device', drive)
548 result = event
549 cancelled = True
550 elif event['event'] == 'JOB_STATUS_CHANGE':
551 self.assert_qmp(event, 'data/id', drive)
554 self.assert_no_active_block_jobs()
555 return result
557 def wait_until_completed(self, drive='drive0', check_offset=True):
558 '''Wait for a block job to finish, returning the event'''
559 while True:
560 for event in self.vm.get_qmp_events(wait=True):
561 if event['event'] == 'BLOCK_JOB_COMPLETED':
562 self.assert_qmp(event, 'data/device', drive)
563 self.assert_qmp_absent(event, 'data/error')
564 if check_offset:
565 self.assert_qmp(event, 'data/offset', event['data']['len'])
566 self.assert_no_active_block_jobs()
567 return event
568 elif event['event'] == 'JOB_STATUS_CHANGE':
569 self.assert_qmp(event, 'data/id', drive)
571 def wait_ready(self, drive='drive0'):
572 '''Wait until a block job BLOCK_JOB_READY event'''
573 f = {'data': {'type': 'mirror', 'device': drive } }
574 event = self.vm.event_wait(name='BLOCK_JOB_READY', match=f)
576 def wait_ready_and_cancel(self, drive='drive0'):
577 self.wait_ready(drive=drive)
578 event = self.cancel_and_wait(drive=drive)
579 self.assertEquals(event['event'], 'BLOCK_JOB_COMPLETED')
580 self.assert_qmp(event, 'data/type', 'mirror')
581 self.assert_qmp(event, 'data/offset', event['data']['len'])
583 def complete_and_wait(self, drive='drive0', wait_ready=True):
584 '''Complete a block job and wait for it to finish'''
585 if wait_ready:
586 self.wait_ready(drive=drive)
588 result = self.vm.qmp('block-job-complete', device=drive)
589 self.assert_qmp(result, 'return', {})
591 event = self.wait_until_completed(drive=drive)
592 self.assert_qmp(event, 'data/type', 'mirror')
594 def pause_wait(self, job_id='job0'):
595 with Timeout(1, "Timeout waiting for job to pause"):
596 while True:
597 result = self.vm.qmp('query-block-jobs')
598 found = False
599 for job in result['return']:
600 if job['device'] == job_id:
601 found = True
602 if job['paused'] == True and job['busy'] == False:
603 return job
604 break
605 assert found
607 def pause_job(self, job_id='job0', wait=True):
608 result = self.vm.qmp('block-job-pause', device=job_id)
609 self.assert_qmp(result, 'return', {})
610 if wait:
611 return self.pause_wait(job_id)
612 return result
615 def notrun(reason):
616 '''Skip this test suite'''
617 # Each test in qemu-iotests has a number ("seq")
618 seq = os.path.basename(sys.argv[0])
620 open('%s/%s.notrun' % (output_dir, seq), 'wb').write(reason + '\n')
621 print('%s not run: %s' % (seq, reason))
622 sys.exit(0)
624 def verify_image_format(supported_fmts=[], unsupported_fmts=[]):
625 assert not (supported_fmts and unsupported_fmts)
627 if 'generic' in supported_fmts and \
628 os.environ.get('IMGFMT_GENERIC', 'true') == 'true':
629 # similar to
630 # _supported_fmt generic
631 # for bash tests
632 return
634 not_sup = supported_fmts and (imgfmt not in supported_fmts)
635 if not_sup or (imgfmt in unsupported_fmts):
636 notrun('not suitable for this image format: %s' % imgfmt)
638 def verify_protocol(supported=[], unsupported=[]):
639 assert not (supported and unsupported)
641 if 'generic' in supported:
642 return
644 not_sup = supported and (imgproto not in supported)
645 if not_sup or (imgproto in unsupported):
646 notrun('not suitable for this protocol: %s' % imgproto)
648 def verify_platform(supported_oses=['linux']):
649 if True not in [sys.platform.startswith(x) for x in supported_oses]:
650 notrun('not suitable for this OS: %s' % sys.platform)
652 def verify_cache_mode(supported_cache_modes=[]):
653 if supported_cache_modes and (cachemode not in supported_cache_modes):
654 notrun('not suitable for this cache mode: %s' % cachemode)
656 def supports_quorum():
657 return 'quorum' in qemu_img_pipe('--help')
659 def verify_quorum():
660 '''Skip test suite if quorum support is not available'''
661 if not supports_quorum():
662 notrun('quorum support missing')
664 def main(supported_fmts=[], supported_oses=['linux'], supported_cache_modes=[],
665 unsupported_fmts=[]):
666 '''Run tests'''
668 global debug
670 # We are using TEST_DIR and QEMU_DEFAULT_MACHINE as proxies to
671 # indicate that we're not being run via "check". There may be
672 # other things set up by "check" that individual test cases rely
673 # on.
674 if test_dir is None or qemu_default_machine is None:
675 sys.stderr.write('Please run this test via the "check" script\n')
676 sys.exit(os.EX_USAGE)
678 debug = '-d' in sys.argv
679 verbosity = 1
680 verify_image_format(supported_fmts, unsupported_fmts)
681 verify_platform(supported_oses)
682 verify_cache_mode(supported_cache_modes)
684 # We need to filter out the time taken from the output so that qemu-iotest
685 # can reliably diff the results against master output.
686 import StringIO
687 if debug:
688 output = sys.stdout
689 verbosity = 2
690 sys.argv.remove('-d')
691 else:
692 output = StringIO.StringIO()
694 logging.basicConfig(level=(logging.DEBUG if debug else logging.WARN))
696 class MyTestRunner(unittest.TextTestRunner):
697 def __init__(self, stream=output, descriptions=True, verbosity=verbosity):
698 unittest.TextTestRunner.__init__(self, stream, descriptions, verbosity)
700 # unittest.main() will use sys.exit() so expect a SystemExit exception
701 try:
702 unittest.main(testRunner=MyTestRunner)
703 finally:
704 if not debug:
705 sys.stderr.write(re.sub(r'Ran (\d+) tests? in [\d.]+s', r'Ran \1 tests', output.getvalue()))