unittest_runner.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
# Copyright 2015 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import logging
import os
import sys
import tempfile
import time
import unittest

WORKING_DIR = os.path.dirname(os.path.realpath(__file__))


def PassMessage(message):
  print '\033[22;32m%s\033[22;0m' % message


def FailMessage(message):
  print '\033[22;31m%s\033[22;0m' % message


def LogFilePath(dir_name, filename):
  return os.path.join(dir_name, filename.replace('/', '_') + '.log')


def main(argv):
  logging.disable(logging.CRITICAL)  # Ignore all logging output from unittest.
  temp_dir = tempfile.mkdtemp()
  total_tests = 0
  passed_tests = 0
  failed_tests = []
  for filename in argv[1:]:
    if not os.path.isfile(filename):
      FailMessage('File is not found: %s' % filename)
      continue
    if not filename.endswith('.py'):
      FailMessage('File is not python file: %s' % filename)
      continue
    filename = os.path.relpath(filename, WORKING_DIR)
    module_name = filename[:-3].replace('/', '.')

    suite = unittest.defaultTestLoader.loadTestsFromName(module_name)
    start_time = time.time()
    log_file = LogFilePath(temp_dir, filename)
    result = unittest.TextTestRunner(stream=open(log_file, 'w')).run(suite)
    duration = time.time() - start_time
    if result.wasSuccessful():
      PassMessage('*** PASS [%.2f s] %s' % (duration, filename))
      passed_tests += 1
    else:
      FailMessage('*** FAIL [%.2f s] %s' % (duration, filename))
      failed_tests.append(filename)
    total_tests += 1

  print '%d/%d tests passed' % (passed_tests, total_tests)
  if failed_tests:
    print ''
    print 'Logs of %d failed tests:' % len(failed_tests)
    for failed_test in failed_tests:
      FailMessage(LogFilePath(temp_dir, failed_test))
    print 'To re-test failed unittests, run:'
    print 'make test UNITTEST_WHITELIST="%s"' % ' '.join(failed_tests)


if __name__ == '__main__':
  main(sys.argv)