Bug 1912292 - Run puppeteer tests in several chunks. r=webdriver-reviewers,whimboo,ahal

Differential Revision: https://phabricator.services.mozilla.com/D223500

remote/doc/Testing.md

@@ -149,6 +149,12 @@ test files have to be updated first. To select specific tests or test
 groups within a file define [exclusive tests] by adding the `.only` suffix
 like `it.only()` or `describe.only()`.
 
+It is also possible, similar to how it works in CI, to run tests in chunks:
+
+```shell
+% ./mach puppeteer-test --this-chunk=1 --total-chunks=2
+```
+
 More customizations for [Mocha] can be found in its own documentation.
 
 Test expectation metadata is collected in _remote/test/puppeteer-expected.json_

remote/mach_commands.py

@@ -409,6 +409,8 @@ class PuppeteerRunner(MozbuildObject):
         binary = params.get("binary")
         headless = params.get("headless", False)
         product = params.get("product", "firefox")
+        this_chunk = params.get("this_chunk", "1")
+        total_chunks = params.get("total_chunks", "1")
         with_cdp = params.get("cdp", False)
 
         extra_options = {}
@@ -458,26 +460,34 @@ class PuppeteerRunner(MozbuildObject):
         if product == "chrome":
             if with_cdp:
                 if headless:
-                    test_command = "test:chrome:headless"
+                    test_command = "chrome-headless"
                 else:
-                    test_command = "test:chrome:headful"
+                    test_command = "chrome-headful"
             elif headless:
-                test_command = "test:chrome:bidi"
+                test_command = "chrome-bidi"
             else:
                 raise Exception(
                     "Chrome doesn't support headful mode with the WebDriver BiDi protocol"
                 )
         elif product == "firefox":
             if with_cdp:
-                test_command = "test:firefox:cdp"
+                test_command = "firefox-cdp"
             elif headless:
-                test_command = "test:firefox:headless"
+                test_command = "firefox-headless"
             else:
-                test_command = "test:firefox:headful"
+                test_command = "firefox-headful"
         else:
-            test_command = "test:" + product
+            test_command = product
 
-        command = ["run", test_command, "--"] + mocha_options
+        command = [
+            "run",
+            "test",
+            "--",
+            "--shard",
+            "'" + this_chunk + "-" + total_chunks + "'",
+            "--test-suite",
+            test_command,
+        ] + mocha_options
 
         prefs = {}
         for k, v in params.get("extra_prefs", {}).items():
@@ -586,6 +596,18 @@ def create_parser_puppeteer():
         metavar="<option>=<value>",
         help="Defines additional options for `puppeteer.launch`.",
     )
+    p.add_argument(
+        "--this-chunk",
+        type=str,
+        default="1",
+        help="Defines a current chunk to run.",
+    )
+    p.add_argument(
+        "--total-chunks",
+        type=str,
+        default="1",
+        help="Defines a total amount of chunks to run.",
+    )
     p.add_argument(
         "-v",
         dest="verbosity",
@@ -660,6 +682,8 @@ def puppeteer_test(
     verbosity=0,
     tests=None,
     product="firefox",
+    this_chunk="1",
+    total_chunks="1",
     **kwargs,
 ):
     logger = mozlog.commandline.setup_logging(
@@ -718,6 +742,8 @@ def puppeteer_test(
         "extra_prefs": prefs,
         "product": product,
         "extra_launcher_options": options,
+        "this_chunk": this_chunk,
+        "total_chunks": total_chunks,
     }
     puppeteer = command_context._spawn(PuppeteerRunner)
     try:

taskcluster/docs/attributes.rst

@@ -528,3 +528,11 @@ A list of the test manifests that run in this task.
 lull-schedule
 =============
 Used by performance tasks to schedule them at a specified frequency in a best-effort method. Schedules them when the overall CI load is low for a given platform. Use "w" for weeks, "d" for days, "h" for hours, and "m" for minutes in a string like so to specify the scheduling frequency: 1d, 1w 4h, 2w 4d 1h.
+
+this_chunk
+=============
+Used by source tests to support chunking and specify a current chunk.
+
+total_chunks
+=============
+Used by source tests to support chunking and specify a total amount of chunks.

taskcluster/kinds/source-test/kind.yml

@@ -5,6 +5,7 @@
 loader: taskgraph.loader.transform:loader
 
 transforms:
+    - taskgraph.transforms.chunking
     - gecko_taskgraph.transforms.try_job:transforms
     - gecko_taskgraph.transforms.source_test:transforms
     - taskgraph.transforms.task_context

taskcluster/kinds/source-test/puppeteer.yml

@@ -41,13 +41,19 @@ puppeteer:
             $MOZ_FETCHES_DIR/firefox/firefox --screenshot http://example.org &&
             ./mach puppeteer-test --ci -vv --binary $MOZ_FETCHES_DIR/firefox/firefox --cdp --headless --log-tbpl - --log-errorsummary /builds/worker/pup_errorsummary.json
 
-puppeteer-with-bidi:
+puppeteer-with-bidi-{this_chunk}:
     description: Puppeteer tests against Firefox Webdriver BiDi remote protocol
+    chunk:
+        total-chunks: 2
+        substitution-fields:
+            - name
+            - treeherder
+            - run
     treeherder:
-        symbol: Pup(wd)
+        symbol: Pup(Wd{this_chunk})
     run:  # Bug 1651542: Use screenshot feature to warm-up the font cache before the actual test
         using: run-task
         command: >
             cd $GECKO_PATH/ &&
             $MOZ_FETCHES_DIR/firefox/firefox --screenshot http://example.org &&
-            ./mach puppeteer-test --ci -vv --binary $MOZ_FETCHES_DIR/firefox/firefox --headless --log-tbpl - --log-errorsummary /builds/worker/pup_errorsummary.json
+            ./mach puppeteer-test --ci -vv --binary $MOZ_FETCHES_DIR/firefox/firefox --headless --log-tbpl - --log-errorsummary /builds/worker/pup_errorsummary.json --this-chunk={this_chunk} --total-chunks={total_chunks}