5 files changed, 552 insertions, 0 deletions
diff --git a/taskcluster/gecko_taskgraph/loader/__init__.py b/taskcluster/gecko_taskgraph/loader/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
--- /dev/null
+++ b/taskcluster/gecko_taskgraph/loader/__init__.py
diff --git a/taskcluster/gecko_taskgraph/loader/multi_dep.py b/taskcluster/gecko_taskgraph/loader/multi_dep.py
new file mode 100644
index 0000000000..04fda04415
--- /dev/null
+++ b/taskcluster/gecko_taskgraph/loader/multi_dep.py
@@ -0,0 +1,275 @@
+# 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 taskgraph.task import Task
+from taskgraph.util.schema import Schema
+from voluptuous import Required
+
+from gecko_taskgraph.util.copy_task import copy_task
+
+from ..util.attributes import sorted_unique_list
+
+schema = Schema(
+    {
+        Required("primary-dependency", "primary dependency task"): Task,
+        Required(
+            "dependent-tasks",
+            "dictionary of dependent tasks, keyed by kind",
+        ): {str: Task},
+    }
+)
+
+
+# Define a collection of group_by functions
+GROUP_BY_MAP = {}
+
+
+def group_by(name):
+    def wrapper(func):
+        GROUP_BY_MAP[name] = func
+        return func
+
+    return wrapper
+
+
+def loader(kind, path, config, params, loaded_tasks):
+    """
+    Load tasks based on the jobs dependant kinds, designed for use as
+    multiple-dependent needs.
+
+    Required ``group-by-fn`` is used to define how we coalesce the
+    multiple deps together to pass to transforms, e.g. all kinds specified get
+    collapsed by platform with `platform`
+
+    Optional ``primary-dependency`` (ordered list or string) is used to determine
+    which upstream kind to inherit attrs from. See ``get_primary_dep``.
+
+    The `only-for-build-platforms` kind configuration, if specified, will limit
+    the build platforms for which a job will be created. Alternatively there is
+    'not-for-build-platforms' kind configuration which will be consulted only after
+    'only-for-build-platforms' is checked (if present), and omit any jobs where the
+    build platform matches.
+
+    Optional ``job-template`` kind configuration value, if specified, will be used to
+    pass configuration down to the specified transforms used.
+    """
+    job_template = config.get("job-template")
+
+    for dep_tasks in group_tasks(config, loaded_tasks):
+        job = {"dependent-tasks": dep_tasks}
+        job["primary-dependency"] = get_primary_dep(config, dep_tasks)
+        if job_template:
+            job.update(copy_task(job_template))
+        # copy shipping_product from upstream
+        product = job["primary-dependency"].attributes.get(
+            "shipping_product", job["primary-dependency"].task.get("shipping-product")
+        )
+        if product:
+            job.setdefault("shipping-product", product)
+        job.setdefault("attributes", {})["required_signoffs"] = sorted_unique_list(
+            *(
+                task.attributes.get("required_signoffs", [])
+                for task in dep_tasks.values()
+            )
+        )
+
+        yield job
+
+
+def skip_only_or_not(config, task):
+    """Return True if we should skip this task based on `only_` or `not_` config."""
+    only_platforms = config.get("only-for-build-platforms")
+    not_platforms = config.get("not-for-build-platforms")
+    only_attributes = config.get("only-for-attributes")
+    not_attributes = config.get("not-for-attributes")
+    task_attrs = task.attributes
+    if only_platforms or not_platforms:
+        platform = task_attrs.get("build_platform")
+        build_type = task_attrs.get("build_type")
+        if not platform or not build_type:
+            return True
+        combined_platform = f"{platform}/{build_type}"
+        if only_platforms and combined_platform not in only_platforms:
+            return True
+        if not_platforms and combined_platform in not_platforms:
+            return True
+    if only_attributes:
+        if not set(only_attributes) & set(task_attrs):
+            # make sure any attribute exists
+            return True
+    if not_attributes:
+        if set(not_attributes) & set(task_attrs):
+            return True
+    return False
+
+
+def group_tasks(config, tasks):
+    group_by_fn = GROUP_BY_MAP[config["group-by"]]
+
+    groups = group_by_fn(config, tasks)
+
+    for combinations in groups.values():
+        kinds = [f.kind for f in combinations]
+        assert_unique_members(
+            kinds,
+            error_msg=("Multi_dep.py should have filtered down to one task per kind"),
+        )
+        dependencies = {t.kind: copy_task(t) for t in combinations}
+        yield dependencies
+
+
+@group_by("platform")
+def platform_grouping(config, tasks):
+    groups = {}
+    for task in tasks:
+        if task.kind not in config.get("kind-dependencies", []):
+            continue
+        if skip_only_or_not(config, task):
+            continue
+        platform = task.attributes.get("build_platform")
+        build_type = task.attributes.get("build_type")
+        product = task.attributes.get(
+            "shipping_product", task.task.get("shipping-product")
+        )
+
+        groups.setdefault((platform, build_type, product), []).append(task)
+    return groups
+
+
+@group_by("single-locale")
+def single_locale_grouping(config, tasks):
+    """Split by a single locale (but also by platform, build-type, product)
+
+    The locale can be `None` (en-US build/signing/repackage), a single locale,
+    or multiple locales per task, e.g. for l10n chunking. In the case of a task
+    with, say, five locales, the task will show up in all five locale groupings.
+
+    This grouping is written for non-partner-repack beetmover, but might also
+    be useful elsewhere.
+
+    """
+    groups = {}
+
+    for task in tasks:
+        if task.kind not in config.get("kind-dependencies", []):
+            continue
+        if skip_only_or_not(config, task):
+            continue
+        platform = task.attributes.get("build_platform")
+        build_type = task.attributes.get("build_type")
+        product = task.attributes.get(
+            "shipping_product", task.task.get("shipping-product")
+        )
+        task_locale = task.attributes.get("locale")
+        chunk_locales = task.attributes.get("chunk_locales")
+        locales = chunk_locales or [task_locale]
+
+        for locale in locales:
+            locale_key = (platform, build_type, product, locale)
+            groups.setdefault(locale_key, [])
+            if task not in groups[locale_key]:
+                groups[locale_key].append(task)
+
+    return groups
+
+
+@group_by("chunk-locales")
+def chunk_locale_grouping(config, tasks):
+    """Split by a chunk_locale (but also by platform, build-type, product)
+
+    This grouping is written for mac signing with notarization, but might also
+    be useful elsewhere.
+
+    """
+    groups = {}
+
+    for task in tasks:
+        if task.kind not in config.get("kind-dependencies", []):
+            continue
+        if skip_only_or_not(config, task):
+            continue
+        platform = task.attributes.get("build_platform")
+        build_type = task.attributes.get("build_type")
+        product = task.attributes.get(
+            "shipping_product", task.task.get("shipping-product")
+        )
+        chunk_locales = tuple(sorted(task.attributes.get("chunk_locales", [])))
+
+        chunk_locale_key = (platform, build_type, product, chunk_locales)
+        groups.setdefault(chunk_locale_key, [])
+        if task not in groups[chunk_locale_key]:
+            groups[chunk_locale_key].append(task)
+
+    return groups
+
+
+@group_by("partner-repack-ids")
+def partner_repack_ids_grouping(config, tasks):
+    """Split by partner_repack_ids (but also by platform, build-type, product)
+
+    This grouping is written for release-{eme-free,partner}-repack-signing.
+
+    """
+    groups = {}
+
+    for task in tasks:
+        if task.kind not in config.get("kind-dependencies", []):
+            continue
+        if skip_only_or_not(config, task):
+            continue
+        platform = task.attributes.get("build_platform")
+        build_type = task.attributes.get("build_type")
+        product = task.attributes.get(
+            "shipping_product", task.task.get("shipping-product")
+        )
+        partner_repack_ids = tuple(
+            sorted(task.task.get("extra", {}).get("repack_ids", []))
+        )
+
+        partner_repack_ids_key = (platform, build_type, product, partner_repack_ids)
+        groups.setdefault(partner_repack_ids_key, [])
+        if task not in groups[partner_repack_ids_key]:
+            groups[partner_repack_ids_key].append(task)
+
+    return groups
+
+
+def assert_unique_members(kinds, error_msg=None):
+    if len(kinds) != len(set(kinds)):
+        raise Exception(error_msg)
+
+
+def get_primary_dep(config, dep_tasks):
+    """Find the dependent task to inherit attributes from.
+
+    If ``primary-dependency`` is defined in ``kind.yml`` and is a string,
+    then find the first dep with that task kind and return it. If it is
+    defined and is a list, the first kind in that list with a matching dep
+    is the primary dependency. If it's undefined, return the first dep.
+
+    """
+    primary_dependencies = config.get("primary-dependency")
+    if isinstance(primary_dependencies, str):
+        primary_dependencies = [primary_dependencies]
+    if not primary_dependencies:
+        assert len(dep_tasks) == 1, "Must define a primary-dependency!"
+        return list(dep_tasks.values())[0]
+    primary_dep = None
+    for primary_kind in primary_dependencies:
+        for dep_kind in dep_tasks:
+            if dep_kind == primary_kind:
+                assert (
+                    primary_dep is None
+                ), "Too many primary dependent tasks in dep_tasks: {}!".format(
+                    [t.label for t in dep_tasks]
+                )
+                primary_dep = dep_tasks[dep_kind]
+    if primary_dep is None:
+        raise Exception(
+            "Can't find dependency of {}: {}".format(
+                config["primary-dependency"], config
+            )
+        )
+    return primary_dep
diff --git a/taskcluster/gecko_taskgraph/loader/single_dep.py b/taskcluster/gecko_taskgraph/loader/single_dep.py
new file mode 100644
index 0000000000..caea4b85d0
--- /dev/null
+++ b/taskcluster/gecko_taskgraph/loader/single_dep.py
@@ -0,0 +1,76 @@
+# 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 taskgraph.task import Task
+from taskgraph.util.schema import Schema
+from voluptuous import Required
+
+from gecko_taskgraph.util.copy_task import copy_task
+
+schema = Schema(
+    {
+        Required("primary-dependency", "primary dependency task"): Task,
+    }
+)
+
+
+def loader(kind, path, config, params, loaded_tasks):
+    """
+    Load tasks based on the jobs dependant kinds.
+
+    The `only-for-build-platforms` kind configuration, if specified, will limit
+    the build platforms for which a job will be created. Alternatively there is
+    'not-for-build-platforms' kind configuration which will be consulted only after
+    'only-for-build-platforms' is checked (if present), and omit any jobs where the
+    build platform matches.
+
+    Optional `only-for-attributes` kind configuration, if specified, will limit
+    the jobs chosen to ones which have the specified attribute, with the specified
+    value.
+
+    Optional `job-template` kind configuration value, if specified, will be used to
+    pass configuration down to the specified transforms used.
+    """
+    only_platforms = config.get("only-for-build-platforms")
+    not_platforms = config.get("not-for-build-platforms")
+    only_attributes = config.get("only-for-attributes")
+    job_template = config.get("job-template")
+
+    for task in loaded_tasks:
+        if task.kind not in config.get("kind-dependencies", []):
+            continue
+
+        if only_platforms or not_platforms:
+            build_platform = task.attributes.get("build_platform")
+            build_type = task.attributes.get("build_type")
+            if not build_platform or not build_type:
+                continue
+            platform = f"{build_platform}/{build_type}"
+            if only_platforms and platform not in only_platforms:
+                continue
+            elif not_platforms and platform in not_platforms:
+                continue
+
+        if only_attributes:
+            config_attrs = set(only_attributes)
+            if not config_attrs & set(task.attributes):
+                # make sure any attribute exists
+                continue
+
+        job = {
+            "primary-dependency": task,
+        }
+
+        if job_template:
+            job.update(copy_task(job_template))
+
+        # copy shipping_product from upstream
+        product = task.attributes.get(
+            "shipping_product", task.task.get("shipping-product")
+        )
+        if product:
+            job.setdefault("shipping-product", product)
+
+        yield job
diff --git a/taskcluster/gecko_taskgraph/loader/test.py b/taskcluster/gecko_taskgraph/loader/test.py
new file mode 100644
index 0000000000..c97acecd1a
--- /dev/null
+++ b/taskcluster/gecko_taskgraph/loader/test.py
@@ -0,0 +1,142 @@
+# 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 logging
+
+from taskgraph.util.yaml import load_yaml
+
+from gecko_taskgraph.util.copy_task import copy_task
+
+from .transform import loader as transform_loader
+
+logger = logging.getLogger(__name__)
+
+
+def loader(kind, path, config, params, loaded_tasks):
+    """
+    Generate tasks implementing Gecko tests.
+    """
+
+    builds_by_platform = get_builds_by_platform(
+        dep_kind="build", loaded_tasks=loaded_tasks
+    )
+    signed_builds_by_platform = get_builds_by_platform(
+        dep_kind="build-signing", loaded_tasks=loaded_tasks
+    )
+
+    # get the test platforms for those build tasks
+    test_platforms_cfg = load_yaml(path, "test-platforms.yml")
+    test_platforms = get_test_platforms(
+        test_platforms_cfg, builds_by_platform, signed_builds_by_platform
+    )
+
+    # expand the test sets for each of those platforms
+    test_sets_cfg = load_yaml(path, "test-sets.yml")
+    test_platforms = expand_tests(test_sets_cfg, test_platforms)
+
+    # load the test descriptions
+    tests = transform_loader(kind, path, config, params, loaded_tasks)
+    test_descriptions = {t.pop("name"): t for t in tests}
+
+    # generate all tests for all test platforms
+    for test_platform_name, test_platform in test_platforms.items():
+        for test_name in test_platform["test-names"]:
+            test = copy_task(test_descriptions[test_name])
+            test["build-platform"] = test_platform["build-platform"]
+            test["test-platform"] = test_platform_name
+            test["build-label"] = test_platform["build-label"]
+            if test_platform.get("build-signing-label", None):
+                test["build-signing-label"] = test_platform["build-signing-label"]
+
+            test["build-attributes"] = test_platform["build-attributes"]
+            test["test-name"] = test_name
+            if test_platform.get("shippable"):
+                test.setdefault("attributes", {})["shippable"] = True
+                test["attributes"]["shipping_product"] = test_platform[
+                    "shipping_product"
+                ]
+
+            logger.debug(
+                "Generating tasks for test {} on platform {}".format(
+                    test_name, test["test-platform"]
+                )
+            )
+            yield test
+
+
+def get_builds_by_platform(dep_kind, loaded_tasks):
+    """Find the build tasks on which tests will depend, keyed by
+    platform/type.  Returns a dictionary mapping build platform to task."""
+    builds_by_platform = {}
+    for task in loaded_tasks:
+        if task.kind != dep_kind:
+            continue
+
+        build_platform = task.attributes.get("build_platform")
+        build_type = task.attributes.get("build_type")
+        if not build_platform or not build_type:
+            continue
+        platform = f"{build_platform}/{build_type}"
+        if platform in builds_by_platform:
+            raise Exception("multiple build jobs for " + platform)
+        builds_by_platform[platform] = task
+    return builds_by_platform
+
+
+def get_test_platforms(
+    test_platforms_cfg, builds_by_platform, signed_builds_by_platform={}
+):
+    """Get the test platforms for which test tasks should be generated,
+    based on the available build platforms.  Returns a dictionary mapping
+    test platform to {test-set, build-platform, build-label}."""
+    test_platforms = {}
+    for test_platform, cfg in test_platforms_cfg.items():
+        build_platform = cfg["build-platform"]
+        if build_platform not in builds_by_platform:
+            logger.warning(
+                "No build task with platform {}; ignoring test platform {}".format(
+                    build_platform, test_platform
+                )
+            )
+            continue
+        test_platforms[test_platform] = {
+            "build-platform": build_platform,
+            "build-label": builds_by_platform[build_platform].label,
+            "build-attributes": builds_by_platform[build_platform].attributes,
+        }
+
+        if builds_by_platform[build_platform].attributes.get("shippable"):
+            test_platforms[test_platform]["shippable"] = builds_by_platform[
+                build_platform
+            ].attributes["shippable"]
+            test_platforms[test_platform]["shipping_product"] = builds_by_platform[
+                build_platform
+            ].attributes["shipping_product"]
+
+        test_platforms[test_platform].update(cfg)
+
+    return test_platforms
+
+
+def expand_tests(test_sets_cfg, test_platforms):
+    """Expand the test sets in `test_platforms` out to sets of test names.
+    Returns a dictionary like `get_test_platforms`, with an additional
+    `test-names` key for each test platform, containing a set of test
+    names."""
+    rv = {}
+    for test_platform, cfg in test_platforms.items():
+        test_sets = cfg["test-sets"]
+        if not set(test_sets) <= set(test_sets_cfg):
+            raise Exception(
+                "Test sets {} for test platform {} are not defined".format(
+                    ", ".join(test_sets), test_platform
+                )
+            )
+        test_names = set()
+        for test_set in test_sets:
+            test_names.update(test_sets_cfg[test_set])
+        rv[test_platform] = cfg.copy()
+        rv[test_platform]["test-names"] = test_names
+    return rv
diff --git a/taskcluster/gecko_taskgraph/loader/transform.py b/taskcluster/gecko_taskgraph/loader/transform.py
new file mode 100644
index 0000000000..1e513bcb73
--- /dev/null
+++ b/taskcluster/gecko_taskgraph/loader/transform.py
@@ -0,0 +1,59 @@
+# 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 logging
+
+from taskgraph.util.yaml import load_yaml
+
+from ..util.templates import merge
+
+logger = logging.getLogger(__name__)
+
+
+def loader(kind, path, config, params, loaded_tasks):
+    """
+    Get the input elements that will be transformed into tasks in a generic
+    way.  The elements themselves are free-form, and become the input to the
+    first transform.
+
+    By default, this reads jobs from the `jobs` key, or from yaml files
+    named by `jobs-from`.  The entities are read from mappings, and the
+    keys to those mappings are added in the `name` key of each entity.
+
+    If there is a `job-defaults` config, then every job is merged with it.
+    This provides a simple way to set default values for all jobs of a kind.
+    The `job-defaults` key can also be specified in a yaml file pointed to by
+    `jobs-from`. In this case it will only apply to tasks defined in the same
+    file.
+
+    Other kind implementations can use a different loader function to
+    produce inputs and hand them to `transform_inputs`.
+    """
+
+    def jobs():
+        defaults = config.get("job-defaults")
+        for name, job in config.get("jobs", {}).items():
+            if defaults:
+                job = merge(defaults, job)
+            job["job-from"] = "kind.yml"
+            yield name, job
+
+        for filename in config.get("jobs-from", []):
+            tasks = load_yaml(path, filename)
+
+            file_defaults = tasks.pop("job-defaults", None)
+            if defaults:
+                file_defaults = merge(defaults, file_defaults or {})
+
+            for name, job in tasks.items():
+                if file_defaults:
+                    job = merge(file_defaults, job)
+                job["job-from"] = filename
+                yield name, job
+
+    for name, job in jobs():
+        job["name"] = name
+        logger.debug(f"Generating tasks for {kind} {name}")
+        yield job