mach_commands.py (11132B)
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 import argparse 6 import logging 7 import os 8 import subprocess 9 import tempfile 10 from concurrent.futures import ThreadPoolExecutor, as_completed, thread 11 12 import mozinfo 13 from mach.decorators import Command, CommandArgument 14 from manifestparser import TestManifest 15 from manifestparser import filters as mpf 16 from mozbuild.util import cpu_count 17 from mozfile import which 18 from tqdm import tqdm 19 20 21 @Command("python", category="devenv", description="Run Python.") 22 @CommandArgument( 23 "--exec-file", default=None, help="Execute this Python file using `exec`" 24 ) 25 @CommandArgument( 26 "--ipython", 27 action="store_true", 28 default=False, 29 help="Use ipython instead of the default Python REPL.", 30 ) 31 @CommandArgument( 32 "--virtualenv", 33 default=None, 34 help="Prepare and use the virtualenv with the provided name. If not specified, " 35 "then the Mach context is used instead.", 36 ) 37 @CommandArgument("args", nargs=argparse.REMAINDER) 38 def python( 39 command_context, 40 exec_file, 41 ipython, 42 virtualenv, 43 args, 44 ): 45 # Avoid logging the command 46 command_context.log_manager.terminal_handler.setLevel(logging.CRITICAL) 47 48 # Note: subprocess requires native strings in os.environ on Windows. 49 append_env = {"PYTHONDONTWRITEBYTECODE": "1"} 50 51 if virtualenv: 52 command_context._virtualenv_name = virtualenv 53 54 if exec_file: 55 command_context.activate_virtualenv() 56 exec(open(exec_file).read()) 57 return 0 58 59 if ipython: 60 if virtualenv: 61 command_context.virtualenv_manager.ensure() 62 python_path = which( 63 "ipython", path=command_context.virtualenv_manager.bin_path 64 ) 65 if not python_path: 66 raise Exception( 67 "--ipython was specified, but the provided " 68 '--virtualenv doesn\'t have "ipython" installed.' 69 ) 70 else: 71 command_context._virtualenv_name = "ipython" 72 command_context.virtualenv_manager.ensure() 73 python_path = which( 74 "ipython", path=command_context.virtualenv_manager.bin_path 75 ) 76 else: 77 command_context.virtualenv_manager.ensure() 78 python_path = command_context.virtualenv_manager.python_path 79 80 return command_context.run_process( 81 [python_path] + args, 82 pass_thru=True, # Allow user to run Python interactively. 83 ensure_exit_code=False, # Don't throw on non-zero exit code. 84 python_unbuffered=False, # Leave input buffered. 85 append_env=append_env, 86 ) 87 88 89 @Command( 90 "python-test", 91 category="testing", 92 virtualenv_name="python-test", 93 description="Run Python unit tests with pytest.", 94 ) 95 @CommandArgument( 96 "-v", "--verbose", default=False, action="store_true", help="Verbose output." 97 ) 98 @CommandArgument( 99 "-j", 100 "--jobs", 101 default=None, 102 type=int, 103 help="Number of concurrent jobs to run. Default is the number of CPUs " 104 "in the system.", 105 ) 106 @CommandArgument( 107 "-x", 108 "--exitfirst", 109 default=False, 110 action="store_true", 111 help="Runs all tests sequentially and breaks at the first failure.", 112 ) 113 @CommandArgument( 114 "--subsuite", 115 default=None, 116 help=( 117 "Python subsuite to run. If not specified, all subsuites are run. " 118 "Use the string `default` to only run tests without a subsuite." 119 ), 120 ) 121 @CommandArgument( 122 "tests", 123 nargs="*", 124 metavar="TEST", 125 help=( 126 "Tests to run. Each test can be a single file or a directory. " 127 "Default test resolution relies on PYTHON_UNITTEST_MANIFESTS." 128 ), 129 ) 130 @CommandArgument( 131 "extra", 132 nargs=argparse.REMAINDER, 133 metavar="PYTEST ARGS", 134 help=( 135 "Arguments that aren't recognized by mach. These will be " 136 "passed as it is to pytest" 137 ), 138 ) 139 def python_test(command_context, *args, **kwargs): 140 try: 141 tempdir = str(tempfile.mkdtemp(suffix="-python-test")) 142 os.environ["PYTHON_TEST_TMP"] = tempdir 143 return run_python_tests(command_context, *args, **kwargs) 144 finally: 145 import mozfile 146 147 mozfile.remove(tempdir) 148 149 150 def run_python_tests( 151 command_context, 152 tests=None, 153 test_objects=None, 154 subsuite=None, 155 verbose=False, 156 jobs=None, 157 exitfirst=False, 158 extra=None, 159 **kwargs, 160 ): 161 if test_objects is None: 162 from moztest.resolve import TestResolver 163 164 resolver = command_context._spawn(TestResolver) 165 # If we were given test paths, try to find tests matching them. 166 test_objects = resolver.resolve_tests(paths=tests, flavor="python") 167 else: 168 # We've received test_objects from |mach test|. We need to ignore 169 # the subsuite because python-tests don't use this key like other 170 # harnesses do and |mach test| doesn't realize this. 171 subsuite = None 172 173 mp = TestManifest() 174 mp.tests.extend(test_objects) 175 176 filters = [] 177 if subsuite == "default": 178 filters.append(mpf.subsuite(None)) 179 elif subsuite: 180 filters.append(mpf.subsuite(subsuite)) 181 182 tests = mp.active_tests(filters=filters, disabled=False, python=3, **mozinfo.info) 183 184 if not tests: 185 submsg = f"for subsuite '{subsuite}' " if subsuite else "" 186 message = ( 187 "TEST-UNEXPECTED-FAIL | No tests collected " 188 + f"{submsg}(Not in PYTHON_UNITTEST_MANIFESTS?)" 189 ) 190 command_context.log(logging.WARN, "python-test", {}, message) 191 return 1 192 193 parallel = [] 194 sequential = [] 195 os.environ.setdefault("PYTEST_ADDOPTS", "") 196 197 if extra: 198 os.environ["PYTEST_ADDOPTS"] += " " + " ".join(extra) 199 200 installed_requirements = set() 201 for test in tests: 202 if ( 203 test.get("requirements") 204 and test["requirements"] not in installed_requirements 205 ): 206 command_context.virtualenv_manager.install_pip_requirements( 207 test["requirements"], quiet=True 208 ) 209 installed_requirements.add(test["requirements"]) 210 211 if exitfirst: 212 sequential = tests 213 os.environ["PYTEST_ADDOPTS"] += " -x" 214 else: 215 for test in tests: 216 if test.get("sequential"): 217 sequential.append(test) 218 else: 219 parallel.append(test) 220 221 jobs = jobs or cpu_count() 222 223 return_code = 0 224 failure_output = [] 225 226 def on_test_finished(result): 227 output, ret, test_path = result 228 229 if ret: 230 # Log the output of failed tests at the end so it's easy to find. 231 failure_output.extend(output) 232 233 if not return_code: 234 command_context.log( 235 logging.ERROR, 236 "python-test", 237 {"test_path": test_path, "ret": ret}, 238 "Setting retcode to {ret} from {test_path}", 239 ) 240 else: 241 for line in output: 242 command_context.log( 243 logging.INFO, "python-test", {"line": line.rstrip()}, "{line}" 244 ) 245 246 return return_code or ret 247 248 with tqdm( 249 total=(len(parallel) + len(sequential)), 250 unit="Test", 251 desc="Tests Completed", 252 initial=0, 253 ) as progress_bar: 254 try: 255 with ThreadPoolExecutor(max_workers=jobs) as executor: 256 futures = [] 257 258 for test in parallel: 259 command_context.log( 260 logging.DEBUG, 261 "python-test", 262 {"line": f"Launching thread for test {test['file_relpath']}"}, 263 "{line}", 264 ) 265 futures.append( 266 executor.submit( 267 _run_python_test, command_context, test, jobs, verbose 268 ) 269 ) 270 271 try: 272 for future in as_completed(futures): 273 progress_bar.clear() 274 return_code = on_test_finished(future.result()) 275 progress_bar.update(1) 276 except KeyboardInterrupt: 277 # Hack to force stop currently running threads. 278 # https://gist.github.com/clchiou/f2608cbe54403edb0b13 279 executor._threads.clear() 280 thread._threads_queues.clear() 281 raise 282 283 for test in sequential: 284 test_result = _run_python_test(command_context, test, jobs, verbose) 285 286 progress_bar.clear() 287 return_code = on_test_finished(test_result) 288 if return_code and exitfirst: 289 break 290 291 progress_bar.update(1) 292 finally: 293 progress_bar.clear() 294 # Now log all failures (even if there was a KeyboardInterrupt or other exception). 295 for line in failure_output: 296 command_context.log( 297 logging.INFO, "python-test", {"line": line.rstrip()}, "{line}" 298 ) 299 300 command_context.log( 301 logging.INFO, 302 "python-test", 303 {"return_code": return_code}, 304 "Return code from mach python-test: {return_code}", 305 ) 306 307 return return_code 308 309 310 def _run_python_test(command_context, test, jobs, verbose): 311 output = [] 312 313 def _log(line): 314 # Buffer messages if more than one worker to avoid interleaving 315 if jobs > 1: 316 output.append(line) 317 else: 318 command_context.log( 319 logging.INFO, "python-test", {"line": line.rstrip()}, "{line}" 320 ) 321 322 _log(test["path"]) 323 python = command_context.virtualenv_manager.python_path 324 cmd = [python, test["path"]] 325 env = os.environ.copy() 326 env["PYTHONDONTWRITEBYTECODE"] = "1" 327 328 result = subprocess.run( 329 cmd, 330 check=False, 331 env=env, 332 stdout=subprocess.PIPE, 333 stderr=subprocess.STDOUT, 334 text=True, 335 encoding="UTF-8", 336 ) 337 338 return_code = result.returncode 339 340 file_displayed_test = False 341 342 for line in result.stdout.split(os.linesep): 343 if not file_displayed_test: 344 test_ran = "Ran" in line or "collected" in line or line.startswith("TEST-") 345 if test_ran: 346 file_displayed_test = True 347 348 # Hack to make sure treeherder highlights pytest failures 349 if "FAILED" in line.rsplit(" ", 1)[-1]: 350 line = line.replace("FAILED", "TEST-UNEXPECTED-FAIL") 351 352 _log(line) 353 354 if not file_displayed_test: 355 return_code = 1 356 _log( 357 "TEST-UNEXPECTED-FAIL | No test output (missing mozunit.main() " 358 "call?): {}".format(test["path"]) 359 ) 360 361 if verbose: 362 if return_code != 0: 363 _log("Test failed: {}".format(test["path"])) 364 else: 365 _log("Test passed: {}".format(test["path"])) 366 367 return output, return_code, test["path"]