1 files changed, 91 insertions, 0 deletions
diff --git a/taskcluster/gecko_taskgraph/transforms/test/variant.py b/taskcluster/gecko_taskgraph/transforms/test/variant.py
new file mode 100644
index 0000000000..5ed5c52613
--- /dev/null
+++ b/taskcluster/gecko_taskgraph/transforms/test/variant.py
@@ -0,0 +1,91 @@
+# 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/.
+import copy
+
+import jsone
+from taskgraph.transforms.base import TransformSequence
+from taskgraph.util.schema import Schema, validate_schema
+from taskgraph.util.treeherder import join_symbol, split_symbol
+from taskgraph.util.yaml import load_yaml
+from voluptuous import Any, Optional, Required
+
+import gecko_taskgraph
+from gecko_taskgraph.util.templates import merge
+
+transforms = TransformSequence()
+
+TEST_VARIANTS = load_yaml(
+    gecko_taskgraph.GECKO, "taskcluster", "ci", "test", "variants.yml"
+)
+"""List of available test variants defined."""
+
+
+variant_description_schema = Schema(
+    {
+        str: {
+            Required("description"): str,
+            Required("suffix"): str,
+            Required("component"): str,
+            Required("expiration"): str,
+            Optional("when"): {Any("$eval", "$if"): str},
+            Optional("replace"): {str: object},
+            Optional("merge"): {str: object},
+        }
+    }
+)
+"""variant description schema"""
+
+
+@transforms.add
+def split_variants(config, tasks):
+    """Splits test definitions into multiple tasks based on the `variants` key.
+
+    If `variants` are defined, the original task will be yielded along with a
+    copy of the original task for each variant defined in the list. The copies
+    will have the 'unittest_variant' attribute set.
+    """
+    validate_schema(variant_description_schema, TEST_VARIANTS, "In variants.yml:")
+
+    def apply_variant(variant, task):
+        task["description"] = variant["description"].format(**task)
+
+        suffix = f"-{variant['suffix']}"
+        group, symbol = split_symbol(task["treeherder-symbol"])
+        if group != "?":
+            group += suffix
+        else:
+            symbol += suffix
+        task["treeherder-symbol"] = join_symbol(group, symbol)
+
+        # This will be used to set the label and try-name in 'make_job_description'.
+        task.setdefault("variant-suffix", "")
+        task["variant-suffix"] += suffix
+
+        # Replace and/or merge the configuration.
+        task.update(variant.get("replace", {}))
+        return merge(task, variant.get("merge", {}))
+
+    for task in tasks:
+        variants = task.pop("variants", [])
+
+        if task.pop("run-without-variant"):
+            yield copy.deepcopy(task)
+
+        for name in variants:
+            # Apply composite variants (joined by '+') in order.
+            parts = name.split("+")
+            taskv = copy.deepcopy(task)
+            for part in parts:
+                variant = TEST_VARIANTS[part]
+
+                # If any variant in a composite fails this check we skip it.
+                if "when" in variant:
+                    context = {"task": task}
+                    if not jsone.render(variant["when"], context):
+                        break
+
+                taskv = apply_variant(variant, taskv)
+            else:
+                taskv["attributes"]["unittest_variant"] = name
+                yield taskv