From be7de4610e7d55a8b8ef5683e98c4dddd0f16b41 Mon Sep 17 00:00:00 2001
From: Ben Hearsum <ben@mozilla.com>
Date: Fri, 6 Dec 2024 14:37:51 -0500
Subject: [PATCH 1/2] feat: allow task group to be overridden in
 create_tasks/taskgraph_decision

I'm not sure yet if this is a good idea, but it's a possible path forward for https://github.com/mozilla-releng/fxci-config/pull/192, where I'm running an action task that I'd like to have status reported back into the task group that ran the action.
---
 src/taskgraph/create.py   |  4 ++--
 src/taskgraph/decision.py |  3 ++-
 test/test_create.py       | 25 +++++++++++++++++++++++++
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/src/taskgraph/create.py b/src/taskgraph/create.py
index ac026e10d..f76bbbb70 100644
--- a/src/taskgraph/create.py
+++ b/src/taskgraph/create.py
@@ -21,7 +21,7 @@
 testing = False
 
 
-def create_tasks(graph_config, taskgraph, label_to_taskid, params, decision_task_id):
+def create_tasks(graph_config, taskgraph, label_to_taskid, params, decision_task_id, task_group_id=None):
     taskid_to_label = {t: l for l, t in label_to_taskid.items()}
 
     # when running as an actual decision task, we use the decision task's
@@ -43,7 +43,7 @@ def create_tasks(graph_config, taskgraph, label_to_taskid, params, decision_task
         if not any(t in taskgraph.tasks for t in task_def.get("dependencies", [])):
             task_def.setdefault("dependencies", []).append(decision_task_id)
 
-        task_def["taskGroupId"] = decision_task_id
+        task_def["taskGroupId"] = task_group_id or decision_task_id
         task_def["schedulerId"] = scheduler_id
 
     # If `testing` is True, then run without parallelization
diff --git a/src/taskgraph/decision.py b/src/taskgraph/decision.py
index 5398bd70d..5ec59e871 100644
--- a/src/taskgraph/decision.py
+++ b/src/taskgraph/decision.py
@@ -63,7 +63,7 @@ def full_task_graph_to_runnable_tasks(full_task_json):
     return runnable_tasks
 
 
-def taskgraph_decision(options, parameters=None):
+def taskgraph_decision(options, parameters=None, task_group_id=None):
     """
     Run the decision task.  This function implements `mach taskgraph decision`,
     and is responsible for
@@ -152,6 +152,7 @@ def taskgraph_decision(options, parameters=None):
         tgg.label_to_taskid,
         tgg.parameters,
         decision_task_id=decision_task_id,
+        task_group_id=task_group_id,
     )
 
 
diff --git a/test/test_create.py b/test/test_create.py
index 1b93c9b14..eb646bac0 100644
--- a/test/test_create.py
+++ b/test/test_create.py
@@ -58,6 +58,31 @@ def test_create_tasks(self):
                     continue
                 self.assertIn(depid, self.created_tasks)
 
+    def test_create_tasks_explicit_group_id(self):
+        tasks = {
+            "tid-a": Task(
+                kind="test", label="a", attributes={}, task={"payload": "hello world"}
+            ),
+            "tid-b": Task(
+                kind="test", label="b", attributes={}, task={"payload": "hello world"}
+            ),
+        }
+        label_to_taskid = {"a": "tid-a", "b": "tid-b"}
+        graph = Graph(nodes={"tid-a", "tid-b"}, edges={("tid-a", "tid-b", "edge")})
+        taskgraph = TaskGraph(tasks, graph)
+
+        create.create_tasks(
+            GRAPH_CONFIG,
+            taskgraph,
+            label_to_taskid,
+            {"level": "4"},
+            decision_task_id="decisiontask",
+            task_group_id="task_group",
+        )
+
+        for _, task in self.created_tasks.items():
+            self.assertEqual(task["taskGroupId"], "task_group")
+
     def test_create_task_without_dependencies(self):
         "a task with no dependencies depends on the decision task"
         tasks = {

From 71b01223221a50b41100b8a7fdbd9d7d48cd45e4 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 6 Dec 2024 20:38:21 +0000
Subject: [PATCH 2/2] style: pre-commit.ci auto fixes [...]

---
 src/taskgraph/create.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/taskgraph/create.py b/src/taskgraph/create.py
index f76bbbb70..e0faa3dc5 100644
--- a/src/taskgraph/create.py
+++ b/src/taskgraph/create.py
@@ -21,7 +21,14 @@
 testing = False
 
 
-def create_tasks(graph_config, taskgraph, label_to_taskid, params, decision_task_id, task_group_id=None):
+def create_tasks(
+    graph_config,
+    taskgraph,
+    label_to_taskid,
+    params,
+    decision_task_id,
+    task_group_id=None,
+):
     taskid_to_label = {t: l for l, t in label_to_taskid.items()}
 
     # when running as an actual decision task, we use the decision task's