refactor: get rid of hardcoded attribute filtering in find_tasks (#225)

taskcluster/fxci_config_taskgraph/transforms/firefoxci_artifact.py

@@ -23,8 +23,12 @@ def make_firefoxci_artifact_tasks(config, tasks):
     for task in tasks:
         # Format: { <task id>: [<artifact name>]}
         tasks_to_create = defaultdict(list)
+        include_attrs = task.pop("include-attrs", {})
+        exclude_attrs = task.pop("exclude-attrs", {})
         for decision_index_path in task.pop("decision-index-paths"):
-            for task_def in find_tasks(decision_index_path):
+            for task_def in find_tasks(
+                decision_index_path, include_attrs, exclude_attrs
+            ):
                 # Add docker images
                 if "image" in task_def["payload"]:
                     image = task_def["payload"]["image"]

taskcluster/fxci_config_taskgraph/transforms/integration_test.py

@@ -204,6 +204,10 @@ def schedule_tasks_at_index(config, tasks):
         return
 
     for task in tasks:
+        include_attrs = task.pop("include-attrs", {})
+        exclude_attrs = task.pop("exclude-attrs", {})
         for decision_index_path in task.pop("decision-index-paths"):
-            for task_def in find_tasks(decision_index_path):
+            for task_def in find_tasks(
+                decision_index_path, include_attrs, exclude_attrs
+            ):
                 yield make_integration_test_description(task_def, task["name"])

taskcluster/fxci_config_taskgraph/util/integration.py

@@ -7,6 +7,7 @@
 
 import requests
 import taskcluster
+from taskgraph.util.attributes import attrmatch
 
 from fxci_config_taskgraph.util.constants import FIREFOXCI_ROOT_URL
 
@@ -18,8 +19,11 @@ def get_taskcluster_client(service: str):
 
 
 @cache
-def find_tasks(decision_index_path: str) -> list[dict[str, Any]]:
-    """Find tasks targeted by the Decision task pointed to by `decision_index_path`."""
+def _fetch_task_graph(decision_index_path: str) -> list[dict[str, Any]]:
+    """Fetch a decision task's `task-graph.json` given by the
+    `decision-index-path`. This is done separately from `find_tasks`
+    because the @cache decorator does not work with the `dict` parameters
+    in `find_tasks`."""
     queue = get_taskcluster_client("queue")
     index = get_taskcluster_client("index")
 
@@ -37,14 +41,29 @@ def find_tasks(decision_index_path: str) -> list[dict[str, Any]]:
     else:
         task_graph = response
 
+    return task_graph.values()
+
+
+def find_tasks(
+    decision_index_path: str,
+    include_attrs: dict[str, list[str]],
+    exclude_attrs: dict[str, list[str]],
+) -> list[dict[str, Any]]:
+    """Find tasks targeted by the Decision task pointed to by `decision_index_path`
+    that match the the included and excluded attributes given.
+    """
     tasks = []
-    for task in task_graph.values():
+
+    for task in _fetch_task_graph(decision_index_path):
         assert isinstance(task, dict)
         attributes = task.get("attributes", {})
-        if attributes.get("unittest_variant") != "os-integration":
-            continue
-
-        if attributes.get("test_platform", "").startswith(("android-hw", "macosx")):
+        excludes = {
+            key: lambda attr: any([attr.startswith(v) for v in values])
+            for key, values in exclude_attrs.items()
+        }
+        if not attrmatch(attributes, **include_attrs) or attrmatch(
+            attributes, **excludes
+        ):
             continue
 
         tasks.append(task["task"])

taskcluster/kinds/firefoxci-artifact/kind.yml

@@ -34,3 +34,10 @@ tasks:
     description: "Fetch artifacts from Firefox-CI for the integration tests"
     decision-index-paths:
       - gecko.v2.mozilla-central.latest.taskgraph.decision-os-integration
+    include-attrs:
+      unittest_variant:
+        - os-integration
+    exclude-attrs:
+      test_platform:
+        - android-hw
+        - macosx

taskcluster/kinds/integration-test/kind.yml

@@ -16,3 +16,10 @@ tasks:
     description: "Run Gecko integration tests"
     decision-index-paths:
       - gecko.v2.mozilla-central.latest.taskgraph.decision-os-integration
+    include-attrs:
+      unittest_variant:
+        - os-integration
+    exclude-attrs:
+      test_platform:
+        - android-hw
+        - macosx

taskcluster/test/test_transforms_firefoxci_artifact.py

@@ -11,7 +11,7 @@
 
 from fxci_config_taskgraph.transforms.firefoxci_artifact import transforms
 from fxci_config_taskgraph.util.constants import FIREFOXCI_ROOT_URL, STAGING_ROOT_URL
-from fxci_config_taskgraph.util.integration import find_tasks
+from fxci_config_taskgraph.util.integration import _fetch_task_graph
 
 
 @pytest.fixture
@@ -47,8 +47,14 @@ def run_test(monkeypatch, run_transform, responses):
         json={"taskId": decision_task_id},
     )
 
-    def inner(task: dict[str, Any]) -> dict[str, Any] | None:
-        find_tasks.cache_clear()
+    def inner(
+        task: dict[str, Any],
+        include_attrs: dict[str, list[str]] = {"unittest_variant": ["os-integration"]},
+        exclude_attrs: dict[str, list[str]] = {
+            "test_platform": ["android-hw", "macosx"]
+        },
+    ) -> dict[str, Any] | None:
+        _fetch_task_graph.cache_clear()
 
         task = merge(deepcopy(base_task), task)
         task_graph = {task_label: task}
@@ -62,6 +68,8 @@ def inner(task: dict[str, Any]) -> dict[str, Any] | None:
         transform_task = {
             "name": "gecko",
             "decision-index-paths": [index],
+            "include-attrs": include_attrs,
+            "exclude-attrs": exclude_attrs,
             "worker": {
                 "env": {
                     "MOZ_ARTIFACT_DIR": "/builds/worker/artifacts",

taskcluster/test/test_transforms_integration_test.py

@@ -11,7 +11,7 @@
 
 from fxci_config_taskgraph.transforms.integration_test import transforms
 from fxci_config_taskgraph.util.constants import FIREFOXCI_ROOT_URL, STAGING_ROOT_URL
-from fxci_config_taskgraph.util.integration import find_tasks
+from fxci_config_taskgraph.util.integration import _fetch_task_graph
 
 
 @pytest.fixture
@@ -51,8 +51,15 @@ def run_test(monkeypatch, run_transform, responses):
     # this comes from the keys in a `kind`. `task_label` is not really
     # an ideal default, but it's the best option we have here, and this value
     # is irrelevant to many tests anyways.
-    def inner(task: dict[str, Any], name: str = task_label) -> dict[str, Any] | None:
-        find_tasks.cache_clear()
+    def inner(
+        task: dict[str, Any],
+        include_attrs: dict[str, list[str]] = {"unittest_variant": ["os-integration"]},
+        exclude_attrs: dict[str, list[str]] = {
+            "test_platform": ["android-hw", "macosx"]
+        },
+        name: str = task_label,
+    ) -> dict[str, Any] | None:
+        _fetch_task_graph.cache_clear()
 
         task = merge(deepcopy(base_task), task)
         task_graph = {task_label: task}
@@ -64,7 +71,13 @@ def inner(task: dict[str, Any], name: str = task_label) -> dict[str, Any] | None
         )
 
         result = run_transform(
-            transforms, {"decision-index-paths": [index], "name": name}
+            transforms,
+            {
+                "decision-index-paths": [index],
+                "include-attrs": include_attrs,
+                "exclude-attrs": exclude_attrs,
+                "name": name,
+            },
         )
         if not result:
             return None
@@ -115,7 +128,9 @@ def test_android_hw_skipped(run_test):
 
 
 def test_basic(run_test):
-    result = run_test({"attributes": {"unittest_variant": "os-integration"}}, "gecko")
+    result = run_test(
+        {"attributes": {"unittest_variant": "os-integration"}}, name="gecko"
+    )
     assert result == {
         "attributes": {"integration": "gecko"},
         "dependencies": {"apply": "tc-admin-apply-staging"},
@@ -147,7 +162,7 @@ def test_docker_image(run_test):
             "attributes": {"unittest_variant": "os-integration"},
             "task": {"payload": {"image": {"taskId": "def"}}},
         },
-        "gecko",
+        name="gecko",
     )
     assert result["dependencies"] == {
         "apply": "tc-admin-apply-staging",
@@ -240,7 +255,7 @@ def test_private_artifact(run_test):
                 }
             },
         },
-        "gecko",
+        name="gecko",
     )
     assert result["dependencies"] == {
         "apply": "tc-admin-apply-staging",