refactor: get rid of hardcoded attribute filtering in find_tasks

taskcluster/fxci_config_taskgraph/transforms/firefoxci_artifact.py

@@ -25,8 +25,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

@@ -198,6 +198,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)

taskcluster/fxci_config_taskgraph/util/integration.py

@@ -19,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 _find_tasks(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")
 
@@ -38,14 +41,32 @@ 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 _find_tasks(decision_index_path):
         assert isinstance(task, dict)
         attributes = task.get("attributes", {})
-        if attributes.get("unittest_variant") != "os-integration":
-            continue
+        skip = False
+        for attr, values in dict(include_attrs).items():
+            if not any([attributes.get(attr, "").startswith(v) for v in values]):
+                skip = True
+
+        for attr, values in dict(exclude_attrs).items():
+            if any([attributes.get(attr, "").startswith(v) for v in values]):
+                skip = True
 
-        if attributes.get("test_platform", "").startswith(("android-hw", "macosx")):
+        if skip:
             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

@@ -13,7 +13,7 @@
 from fxci_config_taskgraph.util.integration import (
     FIREFOXCI_ROOT_URL,
     STAGING_ROOT_URL,
-    find_tasks,
+    _find_tasks,
 )
 
 
@@ -50,8 +50,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:
+        _find_tasks.cache_clear()
 
         task = merge(deepcopy(base_task), task)
         task_graph = {task_label: task}
@@ -65,6 +71,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

@@ -13,7 +13,7 @@
 from fxci_config_taskgraph.util.integration import (
     FIREFOXCI_ROOT_URL,
     STAGING_ROOT_URL,
-    find_tasks,
+    _find_tasks,
 )
 
 
@@ -50,8 +50,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:
+        _find_tasks.cache_clear()
 
         task = merge(deepcopy(base_task), task)
         task_graph = {task_label: task}
@@ -62,7 +68,14 @@ def inner(task: dict[str, Any]) -> dict[str, Any] | None:
             json=task_graph,
         )
 
-        result = run_transform(transforms, {"decision-index-paths": [index]})
+        result = run_transform(
+            transforms,
+            {
+                "decision-index-paths": [index],
+                "include-attrs": include_attrs,
+                "exclude-attrs": exclude_attrs,
+            },
+        )
         if not result:
             return None
 
@@ -112,7 +125,8 @@ def test_android_hw_skipped(run_test):
 
 
 def test_basic(run_test):
-    result = run_test({"attributes": {"unittest_variant": "os-integration"}})
+    result = run_test(
+        {"attributes": {"unittest_variant": "os-integration"}})
     assert result == {
         "attributes": {"integration": "gecko"},
         "dependencies": {"apply": "tc-admin-apply-staging"},
@@ -136,7 +150,7 @@ def test_docker_image(run_test):
         {
             "attributes": {"unittest_variant": "os-integration"},
             "task": {"payload": {"image": {"taskId": "def"}}},
-        }
+        },
     )
     assert result["dependencies"] == {
         "apply": "tc-admin-apply-staging",
@@ -228,7 +242,7 @@ def test_private_artifact(run_test):
                     "env": {"MOZ_FETCHES": '[{"task": "def", "artifact": "foo.txt"}]'},
                 }
             },
-        }
+        },
     )
     assert result["dependencies"] == {
         "apply": "tc-admin-apply-staging",