Bug 999450 - Add find-test-chunk command in mach to discover the chunk for a mochitest on a platform. r=chmanchester
authorVaibhav Agrawal <vaibhavmagarwal@gmail.com>
Wed, 26 Aug 2015 16:51:15 -0700
changeset 280935 4843e223c10e0614a89602eb0f2a4b1f237d4122
parent 280934 087ba64d44edebd9fe5021384c006ce281f0b0d8
child 280936 4085ce7f545953258bad1e12c6f146ee381b2215
push id8456
push userraliiev@mozilla.com
push dateMon, 21 Sep 2015 14:31:52 +0000
treeherdermozilla-aurora@7f2f0fb041b1 [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewerschmanchester
bugs999450
milestone43.0a1
Bug 999450 - Add find-test-chunk command in mach to discover the chunk for a mochitest on a platform. r=chmanchester
testing/mach_commands.py
testing/mochitest/mach_commands.py
--- a/testing/mach_commands.py
+++ b/testing/mach_commands.py
@@ -1,24 +1,27 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 from __future__ import absolute_import, print_function, unicode_literals
 
+import json
 import os
 import sys
+import tempfile
 
 from mach.decorators import (
     CommandArgument,
     CommandProvider,
     Command,
 )
 
 from mozbuild.base import MachCommandBase
+from argparse import ArgumentParser
 
 
 UNKNOWN_TEST = '''
 I was unable to find tests in the argument(s) given.
 
 You need to specify a test directory, filename, test suite name, or
 abbreviation.
 
@@ -495,8 +498,95 @@ class PushToTry(MachCommandBase):
 
         if verbose:
             print('The following try syntax was calculated:\n\n\t%s\n' % msg)
 
         if push:
             at.push_to_try(msg, verbose)
 
         return
+
+
+def get_parser(argv=None):
+    parser = ArgumentParser()
+    parser.add_argument(dest="suite_name",
+                        nargs=1,
+                        choices=['mochitest'],
+                        type=str,
+                        help="The test for which chunk should be found. It corresponds "
+                             "to the mach test invoked (only 'mochitest' currently).")
+
+    parser.add_argument(dest="test_path",
+                        nargs=1,
+                        type=str,
+                        help="The test (any mochitest) for which chunk should be found.")
+
+    parser.add_argument('--total-chunks',
+                        type=int,
+                        dest='total_chunks',
+                        required=True,
+                        help='Total number of chunks to split tests into.',
+                        default=None
+                        )
+
+    parser.add_argument('-f', "--flavor",
+                        dest="flavor",
+                        type=str,
+                        help="Flavor to which the test belongs to.")
+
+    parser.add_argument('--chunk-by-runtime',
+                        action='store_true',
+                        dest='chunk_by_runtime',
+                        help='Group tests such that each chunk has roughly the same runtime.',
+                        default=False,
+                        )
+
+    parser.add_argument('--chunk-by-dir',
+                        type=int,
+                        dest='chunk_by_dir',
+                        help='Group tests together in the same chunk that are in the same top '
+                             'chunkByDir directories.',
+                        default=None,
+                        )
+
+    return parser
+
+
+@CommandProvider
+class ChunkFinder(MachCommandBase):
+    @Command('find-test-chunk', category='testing',
+             description='Find which chunk a test belongs to (works for mochitest).',
+             parser=get_parser)
+    def chunk_finder(self, **kwargs):
+        flavor = kwargs['flavor']
+        total_chunks = kwargs['total_chunks']
+        test_path = kwargs['test_path'][0]
+        suite_name = kwargs['suite_name'][0]
+        _, dump_tests = tempfile.mkstemp()
+        args = {
+            'totalChunks': total_chunks,
+            'dump_tests': dump_tests,
+            'chunkByDir': kwargs['chunk_by_dir'],
+            'chunkByRuntime': kwargs['chunk_by_runtime'],
+        }
+
+        found = False
+        for this_chunk in range(1, total_chunks+1):
+            args['thisChunk'] = this_chunk
+            try:
+                self._mach_context.commands.dispatch(suite_name, self._mach_context, flavor=flavor, resolve_tests=False, **args)
+            except SystemExit:
+                pass
+            except KeyboardInterrupt:
+                break
+
+            fp = open(os.path.expanduser(args['dump_tests']), 'r')
+            tests = json.loads(fp.read())['active_tests']
+            paths = [t['path'] for t in tests]
+            if test_path in paths:
+                print("The test %s is present in chunk number: %d (it may be skipped)." % (test_path, this_chunk))
+                found = True
+                break
+
+        if not found:
+            raise Exception("Test %s not found." % test_path)
+        # Clean up the file
+        os.remove(dump_tests)
--- a/testing/mochitest/mach_commands.py
+++ b/testing/mochitest/mach_commands.py
@@ -274,19 +274,20 @@ class MochitestRunner(MozbuildObject):
                 imp.load_module('mochitest', fh, path,
                                 ('.py', 'r', imp.PY_SOURCE))
 
             import mochitest
 
         options = Namespace(**kwargs)
 
         from manifestparser import TestManifest
-        manifest = TestManifest()
-        manifest.tests.extend(tests)
-        options.manifestFile = manifest
+        if tests:
+            manifest = TestManifest()
+            manifest.tests.extend(tests)
+            options.manifestFile = manifest
 
         if options.desktop:
             return mochitest.run_desktop_mochitests(options)
 
         try:
             which.which('adb')
         except which.WhichError:
             # TODO Find adb automatically if it isn't on the path
@@ -330,26 +331,27 @@ class MochitestRunner(MozbuildObject):
             options.xrePath = self.get_webapp_runtime_xre_path()
         elif suite == 'webapprt-chrome':
             options.browserArgs.append("-test-mode")
             if not options.app or options.app == self.get_binary_path():
                 options.app = self.get_webapp_runtime_path()
             options.xrePath = self.get_webapp_runtime_xre_path()
 
         from manifestparser import TestManifest
-        manifest = TestManifest()
-        manifest.tests.extend(tests)
-        options.manifestFile = manifest
+        if tests:
+            manifest = TestManifest()
+            manifest.tests.extend(tests)
+            options.manifestFile = manifest
 
-        # When developing mochitest-plain tests, it's often useful to be able to
-        # refresh the page to pick up modifications. Therefore leave the browser
-        # open if only running a single mochitest-plain test. This behaviour can
-        # be overridden by passing in --keep-open=false.
-        if len(tests) == 1 and options.keep_open is None and suite == 'plain':
-            options.keep_open = True
+            # When developing mochitest-plain tests, it's often useful to be able to
+            # refresh the page to pick up modifications. Therefore leave the browser
+            # open if only running a single mochitest-plain test. This behaviour can
+            # be overridden by passing in --keep-open=false.
+            if len(tests) == 1 and options.keep_open is None and suite == 'plain':
+                options.keep_open = True
 
         # We need this to enable colorization of output.
         self.log_manager.enable_unstructured()
         result = mochitest.run_test_harness(options)
         self.log_manager.disable_unstructured()
         return result
 
     def run_android_test(self, context, tests, suite=None, **kwargs):
@@ -362,19 +364,20 @@ class MochitestRunner(MozbuildObject):
         with open(path, 'r') as fh:
             imp.load_module('runtestsremote', fh, path,
                             ('.py', 'r', imp.PY_SOURCE))
         import runtestsremote
 
         options = Namespace(**kwargs)
 
         from manifestparser import TestManifest
-        manifest = TestManifest()
-        manifest.tests.extend(tests)
-        options.manifestFile = manifest
+        if tests:
+            manifest = TestManifest()
+            manifest.tests.extend(tests)
+            options.manifestFile = manifest
 
         return runtestsremote.run_test_harness(options)
 
     def run_robocop_test(self, context, tests, suite=None, **kwargs):
         host_ret = verify_host_bin()
         if host_ret != 0:
             return host_ret
 
@@ -383,19 +386,20 @@ class MochitestRunner(MozbuildObject):
         with open(path, 'r') as fh:
             imp.load_module('runrobocop', fh, path,
                             ('.py', 'r', imp.PY_SOURCE))
         import runrobocop
 
         options = Namespace(**kwargs)
 
         from manifestparser import TestManifest
-        manifest = TestManifest()
-        manifest.tests.extend(tests)
-        options.manifestFile = manifest
+        if tests:
+            manifest = TestManifest()
+            manifest.tests.extend(tests)
+            options.manifestFile = manifest
 
         return runrobocop.run_test_harness(options)
 
 # parser
 
 def setup_argument_parser():
     build_obj = MozbuildObject.from_environment(cwd=here)
 
@@ -454,17 +458,17 @@ class MachCommands(MachCommandBase):
     @Command('mochitest', category='testing',
              conditions=[is_buildapp_in(*SUPPORTED_APPS)],
              description='Run any flavor of mochitest (integration test).',
              parser=setup_argument_parser)
     @CommandArgument('-f', '--flavor',
                      metavar='{{{}}}'.format(', '.join(CANONICAL_FLAVORS)),
                      choices=SUPPORTED_FLAVORS,
                      help='Only run tests of this flavor.')
-    def run_mochitest_general(self, flavor=None, test_objects=None, **kwargs):
+    def run_mochitest_general(self, flavor=None, test_objects=None, resolve_tests=True, **kwargs):
         buildapp = None
         for app in SUPPORTED_APPS:
             if is_buildapp_in(app)(self):
                 buildapp = app
                 break
 
         flavors = None
         if flavor:
@@ -496,17 +500,19 @@ class MachCommands(MachCommandBase):
                 for tp in test_paths:
                     if mozpath.abspath(tp).startswith(gecko_path):
                         new_paths.append(mozpath.relpath(tp, gecko_path))
                     else:
                         new_paths.append(tp)
                 test_paths = new_paths
 
         mochitest = self._spawn(MochitestRunner)
-        tests = mochitest.resolve_tests(test_paths, test_objects, cwd=self._mach_context.cwd)
+        tests = []
+        if resolve_tests:
+            tests = mochitest.resolve_tests(test_paths, test_objects, cwd=self._mach_context.cwd)
 
         subsuite = kwargs.get('subsuite')
         if subsuite == 'default':
             kwargs['subsuite'] = None
 
         suites = defaultdict(list)
         unsupported = set()
         for test in tests:
@@ -525,16 +531,24 @@ class MachCommands(MachCommandBase):
                     unsupported.add(key)
                     continue
             elif subsuite and test['subsuite'] != subsuite:
                 unsupported.add(key)
                 continue
 
             suites[key].append(test)
 
+        # This is a hack to introduce an option in mach to not send
+        # filtered tests to the mochitest harness. Mochitest harness will read
+        # the master manifest in that case.
+        if not resolve_tests:
+            for flavor in flavors:
+                key = (flavor, kwargs.get('subsuite'))
+                suites[key] = []
+
         if not suites:
             # Make it very clear why no tests were found
             if not unsupported:
                 print(TESTS_NOT_FOUND.format('\n'.join(
                     sorted(list(test_paths or test_objects)))))
                 return 1
 
             msg = []