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 mozfile import which
  22 from manifestparser import TestManifest
  23 from manifestparser import filters as mpf
  24
  25 from mozbuild.base import (
  26     MachCommandBase,
  27 )
  28 from mozbuild.virtualenv import VirtualenvManager
  29
  30 from mach.decorators import (
  31     CommandArgument,
  32     CommandProvider,
  33     Command,
  34 )
  35
  36 here = os.path.abspath(os.path.dirname(__file__))
  37
  38
  39 @CommandProvider
  40 class MachCommands(MachCommandBase):
  41     @Command('python', category='devenv',
  42              description='Run Python.')
  43     @CommandArgument('--no-virtualenv', action='store_true',
  44                      help='Do not set up a virtualenv')
  45     @CommandArgument('--exec-file',
  46                      default=None,
  47                      help='Execute this Python file using `exec`')
  48     @CommandArgument('--ipython',
  49                      action='store_true',
  50                      default=False,
  51                      help='Use ipython instead of the default Python REPL.')
  52     @CommandArgument('args', nargs=argparse.REMAINDER)
  53     def python(self, no_virtualenv, exec_file, ipython, args):
  54         # Avoid logging the command
  55         self.log_manager.terminal_handler.setLevel(logging.CRITICAL)
  56
  57         # Note: subprocess requires native strings in os.environ on Windows.
  58         append_env = {
  59             'PYTHONDONTWRITEBYTECODE': str('1'),
  60         }
  61
  62         if no_virtualenv:
  63             python_path = sys.executable
  64             append_env['PYTHONPATH'] = os.pathsep.join(sys.path)
  65         else:
  66             self._activate_virtualenv()
  67             python_path = self.virtualenv_manager.python_path
  68
  69         if exec_file:
  70             exec(open(exec_file).read())
  71             return 0
  72
  73         if ipython:
  74             bindir = os.path.dirname(python_path)
  75             python_path = which('ipython', path=bindir)
  76             if not python_path:
  77                 if not no_virtualenv:
  78                     # Use `_run_pip` directly rather than `install_pip_package` to bypass
  79                     # `req.check_if_exists()` which may detect a system installed ipython.
  80                     self.virtualenv_manager._run_pip(['install', 'ipython'])
  81                     python_path = which('ipython', path=bindir)
  82
  83                 if not python_path:
  84                     print("error: could not detect or install ipython")
  85                     return 1
  86
  87         return self.run_process([python_path] + args,
  88                                 pass_thru=True,  # Allow user to run Python interactively.
  89                                 ensure_exit_code=False,  # Don't throw on non-zero exit code.
  90                                 python_unbuffered=False,  # Leave input buffered.
  91                                 append_env=append_env)
  92
  93     @Command('python-test', category='testing',
  94              description='Run Python unit tests with an appropriate test runner.')
  95     @CommandArgument('-v', '--verbose',
  96                      default=False,
  97                      action='store_true',
  98                      help='Verbose output.')
  99     @CommandArgument('--python',
 100                      default='2.7',
 101                      help='Version of Python for Pipenv to use. When given a '
 102                           'Python version, Pipenv will automatically scan your '
 103                           'system for a Python that matches that given version.')
 104     @CommandArgument('-j', '--jobs',
 105                      default=None,
 106                      type=int,
 107                      help='Number of concurrent jobs to run. Default is the number of CPUs '
 108                           'in the system.')
 109     @CommandArgument('-x', '--exitfirst',
 110                      default=False,
 111                      action='store_true',
 112                      help='Runs all tests sequentially and breaks at the first failure.')
 113     @CommandArgument('--subsuite',
 114                      default=None,
 115                      help=('Python subsuite to run. If not specified, all subsuites are run. '
 116                            'Use the string `default` to only run tests without a subsuite.'))
 117     @CommandArgument('tests', nargs='*',
 118                      metavar='TEST',
 119                      help=('Tests to run. Each test can be a single file or a directory. '
 120                            'Default test resolution relies on PYTHON_UNITTEST_MANIFESTS.'))
 121     @CommandArgument('extra', nargs=argparse.REMAINDER,
 122                      metavar='PYTEST ARGS',
 123                      help=('Arguments that aren\'t recognized by mach. These will be '
 124                            'passed as it is to pytest'))
 125     def python_test(self, *args, **kwargs):
 126         try:
 127             tempdir = os.environ[b'PYTHON_TEST_TMP'] = str(tempfile.mkdtemp(suffix='-python-test'))
 128             return self.run_python_tests(*args, **kwargs)
 129         finally:
 130             import mozfile
 131             mozfile.remove(tempdir)
 132
 133     def run_python_tests(self,
 134                          tests=None,
 135                          test_objects=None,
 136                          subsuite=None,
 137                          verbose=False,
 138                          jobs=None,
 139                          python=None,
 140                          exitfirst=False,
 141                          extra=None,
 142                          **kwargs):
 143         self._activate_test_virtualenvs(python)
 144
 145         if test_objects is None:
 146             from moztest.resolve import TestResolver
 147             resolver = self._spawn(TestResolver)
 148             # If we were given test paths, try to find tests matching them.
 149             test_objects = resolver.resolve_tests(paths=tests, flavor='python')
 150         else:
 151             # We've received test_objects from |mach test|. We need to ignore
 152             # the subsuite because python-tests don't use this key like other
 153             # harnesses do and |mach test| doesn't realize this.
 154             subsuite = None
 155
 156         mp = TestManifest()
 157         mp.tests.extend(test_objects)
 158
 159         filters = []
 160         if subsuite == 'default':
 161             filters.append(mpf.subsuite(None))
 162         elif subsuite:
 163             filters.append(mpf.subsuite(subsuite))
 164
 165         tests = mp.active_tests(
 166             filters=filters,
 167             disabled=False,
 168             python=self.virtualenv_manager.version_info[0],
 169             **mozinfo.info)
 170
 171         if not tests:
 172             submsg = "for subsuite '{}' ".format(subsuite) if subsuite else ""
 173             message = "TEST-UNEXPECTED-FAIL | No tests collected " + \
 174                       "{}(Not in PYTHON_UNITTEST_MANIFESTS?)".format(submsg)
 175             self.log(logging.WARN, 'python-test', {}, message)
 176             return 1
 177
 178         parallel = []
 179         sequential = []
 180         os.environ.setdefault('PYTEST_ADDOPTS', '')
 181
 182         if extra:
 183             os.environ['PYTEST_ADDOPTS'] += " " + " ".join(extra)
 184
 185         if exitfirst:
 186             sequential = tests
 187             os.environ['PYTEST_ADDOPTS'] += " -x"
 188         else:
 189             for test in tests:
 190                 if test.get('sequential'):
 191                     sequential.append(test)
 192                 else:
 193                     parallel.append(test)
 194
 195         self.jobs = jobs or cpu_count()
 196         self.terminate = False
 197         self.verbose = verbose
 198
 199         return_code = 0
 200
 201         def on_test_finished(result):
 202             output, ret, test_path = result
 203
 204             for line in output:
 205                 self.log(logging.INFO, 'python-test', {'line': line.rstrip()}, '{line}')
 206
 207             if ret and not return_code:
 208                 self.log(logging.ERROR, 'python-test', {'test_path': test_path, 'ret': ret},
 209                          'Setting retcode to {ret} from {test_path}')
 210             return return_code or ret
 211
 212         with ThreadPoolExecutor(max_workers=self.jobs) as executor:
 213             futures = [executor.submit(self._run_python_test, test)
 214                        for test in parallel]
 215
 216             try:
 217                 for future in as_completed(futures):
 218                     return_code = on_test_finished(future.result())
 219             except KeyboardInterrupt:
 220                 # Hack to force stop currently running threads.
 221                 # https://gist.github.com/clchiou/f2608cbe54403edb0b13
 222                 executor._threads.clear()
 223                 thread._threads_queues.clear()
 224                 raise
 225
 226         for test in sequential:
 227             return_code = on_test_finished(self._run_python_test(test))
 228             if return_code and exitfirst:
 229                 break
 230
 231         self.log(logging.INFO, 'python-test', {'return_code': return_code},
 232                  'Return code from mach python-test: {return_code}')
 233         return return_code
 234
 235     def _activate_test_virtualenvs(self, python):
 236         """Make sure the test suite virtualenvs are set up and activated.
 237
 238         Args:
 239             python: Optional python version string we want to run the suite with.
 240                 See the `--python` argument to the `mach python-test` command.
 241         """
 242         from mozbuild.pythonutil import find_python3_executable
 243
 244         default_manager = self.virtualenv_manager
 245
 246         # Grab the default virtualenv properties before we activate other virtualenvs.
 247         python = python or default_manager.python_path
 248         py3_root = default_manager.virtualenv_root + '_py3'
 249
 250         self.activate_pipenv(pipfile=None, populate=True, python=python)
 251
 252         # The current process might be running under Python 2 and the Python 3
 253         # virtualenv will not be set up by mach bootstrap. To avoid problems in tests
 254         # that implicitly depend on the Python 3 virtualenv we ensure the Python 3
 255         # virtualenv is up to date before the tests start.
 256         python3, version = find_python3_executable(min_version='3.5.0')
 257
 258         py3_manager = VirtualenvManager(
 259             default_manager.topsrcdir,
 260             default_manager.topobjdir,
 261             py3_root,
 262             default_manager.log_handle,
 263             default_manager.manifest_path,
 264         )
 265         py3_manager.ensure(python3)
 266
 267     def _run_python_test(self, test):
 268         from mozprocess import ProcessHandler
 269
 270         if test.get('requirements'):
 271             self.virtualenv_manager.install_pip_requirements(test['requirements'], quiet=True)
 272
 273         output = []
 274
 275         def _log(line):
 276             # Buffer messages if more than one worker to avoid interleaving
 277             if self.jobs > 1:
 278                 output.append(line)
 279             else:
 280                 self.log(logging.INFO, 'python-test', {'line': line.rstrip()}, '{line}')
 281
 282         file_displayed_test = []  # used as boolean
 283
 284         def _line_handler(line):
 285             if not file_displayed_test:
 286                 output = ('Ran' in line or 'collected' in line or
 287                           line.startswith('TEST-'))
 288                 if output:
 289                     file_displayed_test.append(True)
 290
 291             # Hack to make sure treeherder highlights pytest failures
 292             if b'FAILED' in line.rsplit(b' ', 1)[-1]:
 293                 line = line.replace(b'FAILED', b'TEST-UNEXPECTED-FAIL')
 294
 295             _log(line)
 296
 297         _log(test['path'])
 298         python = self.virtualenv_manager.python_path
 299         cmd = [python, test['path']]
 300         env = os.environ.copy()
 301         env[b'PYTHONDONTWRITEBYTECODE'] = b'1'
 302
 303         # Homebrew on OS X will change Python's sys.executable to a custom value
 304         # which messes with mach's virtualenv handling code. Override Homebrew's
 305         # changes with the correct sys.executable value.
 306         env[b'PYTHONEXECUTABLE'] = python.encode('utf-8')
 307
 308         proc = ProcessHandler(cmd, env=env, processOutputLine=_line_handler, storeOutput=False)
 309         proc.run()
 310
 311         return_code = proc.wait()
 312
 313         if not file_displayed_test:
 314             _log('TEST-UNEXPECTED-FAIL | No test output (missing mozunit.main() '
 315                  'call?): {}'.format(test['path']))
 316
 317         if self.verbose:
 318             if return_code != 0:
 319                 _log('Test failed: {}'.format(test['path']))
 320             else:
 321                 _log('Test passed: {}'.format(test['path']))
 322
 323         return output, return_code, test['path']