feat: add support for fetching upstream tasks in find_tasks

taskcluster/fxci_config_taskgraph/transforms/integration_test.py

@@ -197,5 +197,6 @@ def schedule_tasks_at_index(config, tasks):
 
     for task in tasks:
         for decision_index_path in task.pop("decision-index-paths"):
-            for task_def in find_tasks(decision_index_path):
+            include_deps = task.pop("include-deps", False)
+            for task_def in find_tasks(decision_index_path, include_deps):
                 yield make_integration_test_description(task_def)

taskcluster/fxci_config_taskgraph/util/integration.py

@@ -3,23 +3,39 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
 from functools import cache
-from typing import Any
+import re
+from typing import Any, Union
 
+import os
 import requests
 import taskcluster
+from taskgraph.util.taskcluster import get_ancestors as taskgraph_get_ancestors
 
 FIREFOXCI_ROOT_URL = "https://firefox-ci-tc.services.mozilla.com"
 STAGING_ROOT_URL = "https://stage.taskcluster.nonprod.cloudops.mozgcp.net"
 
 
+def get_ancestors(task_ids: Union[list[str], str]) -> dict[str, str]:
+    # TODO: this is not ideal, but at the moment we don't have a better way
+    # to ensure that the upstream get_ancestors talks to the correct taskcluster
+    # instance.
+    orig = os.environ["TASKCLUSTER_ROOT_URL"]
+    try:
+        os.environ["TASKCLUSTER_ROOT_URL"] = FIREFOXCI_ROOT_URL
+        ret = taskgraph_get_ancestors(task_ids)
+    finally:
+        os.environ["TASKCLUSTER_ROOT_URL"] = orig
+
+    return ret
+
 @cache
 def get_taskcluster_client(service: str):
     options = {"rootUrl": FIREFOXCI_ROOT_URL}
     return getattr(taskcluster, service.capitalize())(options)
 
 
 @cache
-def find_tasks(decision_index_path: str) -> list[dict[str, Any]]:
+def find_tasks(decision_index_path: str, include_deps: list[re.Pattern] = []) -> list[dict[str, Any]]:
     """Find tasks targeted by the Decision task pointed to by `decision_index_path`."""
     queue = get_taskcluster_client("queue")
     index = get_taskcluster_client("index")
@@ -49,5 +65,11 @@ def find_tasks(decision_index_path: str) -> list[dict[str, Any]]:
             continue
 
         tasks.append(task["task"])
+        # get_ancestors can be expensive; don't run it unless we might actually
+        # use the results.
+        if include_deps:
+            for label, task_id in get_ancestors(task["task_id"]).items():
+                if any([pat.match(label) for pat in include_deps]):
+                    tasks.append(queue.task(task_id))
 
     return tasks