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

@@ -4,8 +4,10 @@
 
 from __future__ import absolute_import, print_function, unicode_literals
 
+import json
 import os
 import sys
+import tempfile
 
 from mach.decorators import (
     CommandArgument,
@@ -14,6 +16,7 @@ from mach.decorators import (
 )
 
 from mozbuild.base import MachCommandBase
+from argparse import ArgumentParser
 
 
 UNKNOWN_TEST = '''
@@ -500,3 +503,90 @@ class PushToTry(MachCommandBase):
             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)

testing/mochitest/mach_commands.py

@@ -279,9 +279,10 @@ class MochitestRunner(MozbuildObject):
         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)
@@ -335,16 +336,17 @@ class MochitestRunner(MozbuildObject):
             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()
@@ -367,9 +369,10 @@ class MochitestRunner(MozbuildObject):
         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)
 
@@ -388,9 +391,10 @@ class MochitestRunner(MozbuildObject):
         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)
 
@@ -459,7 +463,7 @@ class MachCommands(MachCommandBase):
                      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):
@@ -501,7 +505,9 @@ class MachCommands(MachCommandBase):
                 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':
@@ -530,6 +536,14 @@ class MachCommands(MachCommandBase):
 
             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: