python/mach_commands.py

   1 # This Source Code Form is subject to the terms of the Mozilla Public
   2 # License, v. 2.0. If a copy of the MPL was not distributed with this
   3 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
   4
   5 from __future__ import absolute_import, print_function, unicode_literals
   6
   7 import argparse
   8 import logging
   9 import os
  10 import sys
  11 import tempfile
  12 from multiprocessing import cpu_count
  13
  14 from concurrent.futures import (
  15     ThreadPoolExecutor,
  16     as_completed,
  17     thread,
  18 )
  19
  20 import mozinfo
  21 from manifestparser import TestManifest
  22 from manifestparser import filters as mpf
  23
  24 from mozbuild.base import (
  25     MachCommandBase,
  26 )
  27
  28 from mach.decorators import (
  29     CommandArgument,
  30     CommandProvider,
  31     Command,
  32 )
  33
  34 here = os.path.abspath(os.path.dirname(__file__))
  35
  36
  37 @CommandProvider
  38 class MachCommands(MachCommandBase):
  39     @Command('python', category='devenv',
  40              description='Run Python.')
  41     @CommandArgument('--no-virtualenv', action='store_true',
  42                      help='Do not set up a virtualenv')
  43     @CommandArgument('--exec-file',
  44                      default=None,
  45                      help='Execute this Python file using `execfile`')
  46     @CommandArgument('args', nargs=argparse.REMAINDER)
  47     def python(self, no_virtualenv, exec_file, args):
  48         # Avoid logging the command
  49         self.log_manager.terminal_handler.setLevel(logging.CRITICAL)
  50
  51         # Note: subprocess requires native strings in os.environ on Windows.
  52         append_env = {
  53             b'PYTHONDONTWRITEBYTECODE': str('1'),
  54         }
  55
  56         if no_virtualenv:
  57             python_path = sys.executable
  58             append_env[b'PYTHONPATH'] = os.pathsep.join(sys.path)
  59         else:
  60             self._activate_virtualenv()
  61             python_path = self.virtualenv_manager.python_path
  62
  63         if exec_file:
  64             execfile(exec_file)
  65             return 0
  66
  67         return self.run_process([python_path] + args,
  68                                 pass_thru=True,  # Allow user to run Python interactively.
  69                                 ensure_exit_code=False,  # Don't throw on non-zero exit code.
  70                                 append_env=append_env)
  71
  72     @Command('python-test', category='testing',
  73              description='Run Python unit tests with an appropriate test runner.')
  74     @CommandArgument('-v', '--verbose',
  75                      default=False,
  76                      action='store_true',
  77                      help='Verbose output.')
  78     @CommandArgument('--python',
  79                      default='2.7',
  80                      help='Version of Python for Pipenv to use. When given a '
  81                           'Python version, Pipenv will automatically scan your '
  82                           'system for a Python that matches that given version.')
  83     @CommandArgument('-j', '--jobs',
  84                      default=None,
  85                      type=int,
  86                      help='Number of concurrent jobs to run. Default is the number of CPUs '
  87                           'in the system.')
  88     @CommandArgument('--subsuite',
  89                      default=None,
  90                      help=('Python subsuite to run. If not specified, all subsuites are run. '
  91                            'Use the string `default` to only run tests without a subsuite.'))
  92     @CommandArgument('tests', nargs='*',
  93                      metavar='TEST',
  94                      help=('Tests to run. Each test can be a single file or a directory. '
  95                            'Default test resolution relies on PYTHON_UNITTEST_MANIFESTS.'))
  96     def python_test(self, *args, **kwargs):
  97         try:
  98             tempdir = os.environ[b'PYTHON_TEST_TMP'] = str(tempfile.mkdtemp(suffix='-python-test'))
  99             return self.run_python_tests(*args, **kwargs)
 100         finally:
 101             import mozfile
 102             mozfile.remove(tempdir)
 103
 104     def run_python_tests(self,
 105                          tests=None,
 106                          test_objects=None,
 107                          subsuite=None,
 108                          verbose=False,
 109                          jobs=None,
 110                          python=None,
 111                          **kwargs):
 112         python = python or self.virtualenv_manager.python_path
 113         self.activate_pipenv(pipfile=None, populate=True, python=python)
 114
 115         if test_objects is None:
 116             from moztest.resolve import TestResolver
 117             resolver = self._spawn(TestResolver)
 118             # If we were given test paths, try to find tests matching them.
 119             test_objects = resolver.resolve_tests(paths=tests, flavor='python')
 120         else:
 121             # We've received test_objects from |mach test|. We need to ignore
 122             # the subsuite because python-tests don't use this key like other
 123             # harnesses do and |mach test| doesn't realize this.
 124             subsuite = None
 125
 126         mp = TestManifest()
 127         mp.tests.extend(test_objects)
 128
 129         filters = []
 130         if subsuite == 'default':
 131             filters.append(mpf.subsuite(None))
 132         elif subsuite:
 133             filters.append(mpf.subsuite(subsuite))
 134
 135         tests = mp.active_tests(
 136             filters=filters,
 137             disabled=False,
 138             python=self.virtualenv_manager.version_info[0],
 139             **mozinfo.info)
 140
 141         if not tests:
 142             submsg = "for subsuite '{}' ".format(subsuite) if subsuite else ""
 143             message = "TEST-UNEXPECTED-FAIL | No tests collected " + \
 144                       "{}(Not in PYTHON_UNITTEST_MANIFESTS?)".format(submsg)
 145             self.log(logging.WARN, 'python-test', {}, message)
 146             return 1
 147
 148         parallel = []
 149         sequential = []
 150         for test in tests:
 151             if test.get('sequential'):
 152                 sequential.append(test)
 153             else:
 154                 parallel.append(test)
 155
 156         self.jobs = jobs or cpu_count()
 157         self.terminate = False
 158         self.verbose = verbose
 159
 160         return_code = 0
 161
 162         def on_test_finished(result):
 163             output, ret, test_path = result
 164
 165             for line in output:
 166                 self.log(logging.INFO, 'python-test', {'line': line.rstrip()}, '{line}')
 167
 168             if ret and not return_code:
 169                 self.log(logging.ERROR, 'python-test', {'test_path': test_path, 'ret': ret},
 170                          'Setting retcode to {ret} from {test_path}')
 171             return return_code or ret
 172
 173         with ThreadPoolExecutor(max_workers=self.jobs) as executor:
 174             futures = [executor.submit(self._run_python_test, test)
 175                        for test in parallel]
 176
 177             try:
 178                 for future in as_completed(futures):
 179                     return_code = on_test_finished(future.result())
 180             except KeyboardInterrupt:
 181                 # Hack to force stop currently running threads.
 182                 # https://gist.github.com/clchiou/f2608cbe54403edb0b13
 183                 executor._threads.clear()
 184                 thread._threads_queues.clear()
 185                 raise
 186
 187         for test in sequential:
 188             return_code = on_test_finished(self._run_python_test(test))
 189
 190         self.log(logging.INFO, 'python-test', {'return_code': return_code},
 191                  'Return code from mach python-test: {return_code}')
 192         return return_code
 193
 194     def _run_python_test(self, test):
 195         from mozprocess import ProcessHandler
 196
 197         if test.get('requirements'):
 198             self.virtualenv_manager.install_pip_requirements(test['requirements'], quiet=True)
 199
 200         output = []
 201
 202         def _log(line):
 203             # Buffer messages if more than one worker to avoid interleaving
 204             if self.jobs > 1:
 205                 output.append(line)
 206             else:
 207                 self.log(logging.INFO, 'python-test', {'line': line.rstrip()}, '{line}')
 208
 209         file_displayed_test = []  # used as boolean
 210
 211         def _line_handler(line):
 212             if not file_displayed_test:
 213                 output = ('Ran' in line or 'collected' in line or
 214                           line.startswith('TEST-'))
 215                 if output:
 216                     file_displayed_test.append(True)
 217
 218             # Hack to make sure treeherder highlights pytest failures
 219             if 'FAILED' in line.rsplit(' ', 1)[-1]:
 220                 line = line.replace('FAILED', 'TEST-UNEXPECTED-FAIL')
 221
 222             _log(line)
 223
 224         _log(test['path'])
 225         cmd = [self.virtualenv_manager.python_path, test['path']]
 226         env = os.environ.copy()
 227         env[b'PYTHONDONTWRITEBYTECODE'] = b'1'
 228
 229         proc = ProcessHandler(cmd, env=env, processOutputLine=_line_handler, storeOutput=False)
 230         proc.run()
 231
 232         return_code = proc.wait()
 233
 234         if not file_displayed_test:
 235             _log('TEST-UNEXPECTED-FAIL | No test output (missing mozunit.main() '
 236                  'call?): {}'.format(test['path']))
 237
 238         if self.verbose:
 239             if return_code != 0:
 240                 _log('Test failed: {}'.format(test['path']))
 241             else:
 242                 _log('Test passed: {}'.format(test['path']))
 243
 244         return output, return_code, test['path']