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         self.activate_pipenv(pipfile=None, populate=True, python=python)
 113
 114         if test_objects is None:
 115             from moztest.resolve import TestResolver
 116             resolver = self._spawn(TestResolver)
 117             # If we were given test paths, try to find tests matching them.
 118             test_objects = resolver.resolve_tests(paths=tests, flavor='python')
 119         else:
 120             # We've received test_objects from |mach test|. We need to ignore
 121             # the subsuite because python-tests don't use this key like other
 122             # harnesses do and |mach test| doesn't realize this.
 123             subsuite = None
 124
 125         mp = TestManifest()
 126         mp.tests.extend(test_objects)
 127
 128         filters = []
 129         if subsuite == 'default':
 130             filters.append(mpf.subsuite(None))
 131         elif subsuite:
 132             filters.append(mpf.subsuite(subsuite))
 133
 134         tests = mp.active_tests(
 135             filters=filters,
 136             disabled=False,
 137             python=self.virtualenv_manager.version_info[0],
 138             **mozinfo.info)
 139
 140         if not tests:
 141             submsg = "for subsuite '{}' ".format(subsuite) if subsuite else ""
 142             message = "TEST-UNEXPECTED-FAIL | No tests collected " + \
 143                       "{}(Not in PYTHON_UNITTEST_MANIFESTS?)".format(submsg)
 144             self.log(logging.WARN, 'python-test', {}, message)
 145             return 1
 146
 147         parallel = []
 148         sequential = []
 149         for test in tests:
 150             if test.get('sequential'):
 151                 sequential.append(test)
 152             else:
 153                 parallel.append(test)
 154
 155         self.jobs = jobs or cpu_count()
 156         self.terminate = False
 157         self.verbose = verbose
 158
 159         return_code = 0
 160
 161         def on_test_finished(result):
 162             output, ret, test_path = result
 163
 164             for line in output:
 165                 self.log(logging.INFO, 'python-test', {'line': line.rstrip()}, '{line}')
 166
 167             if ret and not return_code:
 168                 self.log(logging.ERROR, 'python-test', {'test_path': test_path, 'ret': ret},
 169                          'Setting retcode to {ret} from {test_path}')
 170             return return_code or ret
 171
 172         with ThreadPoolExecutor(max_workers=self.jobs) as executor:
 173             futures = [executor.submit(self._run_python_test, test['path'])
 174                        for test in parallel]
 175
 176             try:
 177                 for future in as_completed(futures):
 178                     return_code = on_test_finished(future.result())
 179             except KeyboardInterrupt:
 180                 # Hack to force stop currently running threads.
 181                 # https://gist.github.com/clchiou/f2608cbe54403edb0b13
 182                 executor._threads.clear()
 183                 thread._threads_queues.clear()
 184                 raise
 185
 186         for test in sequential:
 187             return_code = on_test_finished(self._run_python_test(test['path']))
 188
 189         self.log(logging.INFO, 'python-test', {'return_code': return_code},
 190                  'Return code from mach python-test: {return_code}')
 191         return return_code
 192
 193     def _run_python_test(self, test_path):
 194         from mozprocess import ProcessHandler
 195
 196         output = []
 197
 198         def _log(line):
 199             # Buffer messages if more than one worker to avoid interleaving
 200             if self.jobs > 1:
 201                 output.append(line)
 202             else:
 203                 self.log(logging.INFO, 'python-test', {'line': line.rstrip()}, '{line}')
 204
 205         file_displayed_test = []  # used as boolean
 206
 207         def _line_handler(line):
 208             if not file_displayed_test:
 209                 output = ('Ran' in line or 'collected' in line or
 210                           line.startswith('TEST-'))
 211                 if output:
 212                     file_displayed_test.append(True)
 213
 214             # Hack to make sure treeherder highlights pytest failures
 215             if 'FAILED' in line.rsplit(' ', 1)[-1]:
 216                 line = line.replace('FAILED', 'TEST-UNEXPECTED-FAIL')
 217
 218             _log(line)
 219
 220         _log(test_path)
 221         cmd = [self.virtualenv_manager.python_path, test_path]
 222         env = os.environ.copy()
 223         env[b'PYTHONDONTWRITEBYTECODE'] = b'1'
 224
 225         proc = ProcessHandler(cmd, env=env, processOutputLine=_line_handler, storeOutput=False)
 226         proc.run()
 227
 228         return_code = proc.wait()
 229
 230         if not file_displayed_test:
 231             _log('TEST-UNEXPECTED-FAIL | No test output (missing mozunit.main() '
 232                  'call?): {}'.format(test_path))
 233
 234         if self.verbose:
 235             if return_code != 0:
 236                 _log('Test failed: {}'.format(test_path))
 237             else:
 238                 _log('Test passed: {}'.format(test_path))
 239
 240         return output, return_code, test_path