fix: add pagination for list incomplete tasks

Author: abhishekmadan30

src/taskgraph/util/taskcluster.py

@@ -411,20 +411,27 @@ def send_email(address, subject, content, link):
 def list_task_group_incomplete_tasks(task_group_id):
     """Generate the incomplete tasks in a task group"""
     queue = get_taskcluster_client("queue")
-    response = queue.listTaskGroup(task_group_id)
-
-    if not response or "tasks" not in response:
-        return
-
-    tasks = response.get("tasks", [])
-    for task in tasks:
-        if (status := task.get("status")) is not None:  # type: ignore
-            if (task_id := status.get("taskId")) and status.get("state") in [
-                "running",
-                "pending",
-                "unscheduled",
-            ]:
-                yield task_id
+
+    incomplete_tasks = []
+
+    def pagination_handler(response):
+        if not response or "tasks" not in response:
+            return
+
+        tasks = response.get("tasks", [])
+        for task in tasks:
+            if (status := task.get("status")) is not None:  # type: ignore
+                if (task_id := status.get("taskId")) and status.get("state") in [
+                    "running",
+                    "pending",
+                    "unscheduled",
+                ]:
+                    incomplete_tasks.append(task_id)
+
+    queue.listTaskGroup(task_group_id, paginationHandler=pagination_handler)
+
+    for task_id in incomplete_tasks:
+        yield from task_id
 
 
 @functools.lru_cache(maxsize=None)

test/test_util_taskcluster.py

@@ -412,6 +412,49 @@ def test_list_task_group_incomplete_tasks(responses, root_url):
     assert result == ["1", "2", "3"]
 
 
+def test_list_task_group_incomplete_tasks_with_pagination(responses, root_url):
+    tgid = "abc123"
+    tc.get_taskcluster_client.cache_clear()
+
+    responses.get(
+        f"{root_url}/api/queue/v1/task-group/{tgid}/list",
+        json={
+            "tasks": [
+                {"status": {"taskId": "1", "state": "pending"}},
+                {"status": {"taskId": "2", "state": "completed"}},
+                {"status": {"taskId": "3", "state": "running"}},
+            ],
+            "continuationToken": "page2",
+        },
+    )
+
+    responses.get(
+        f"{root_url}/api/queue/v1/task-group/{tgid}/list",
+        json={
+            "tasks": [
+                {"status": {"taskId": "4", "state": "unscheduled"}},
+                {"status": {"taskId": "5", "state": "failed"}},
+                {"status": {"taskId": "6", "state": "pending"}},
+            ],
+            "continuationToken": "page3",
+        },
+    )
+
+    responses.get(
+        f"{root_url}/api/queue/v1/task-group/{tgid}/list",
+        json={
+            "tasks": [
+                {"status": {"taskId": "7", "state": "running"}},
+                {"status": {"taskId": "8", "state": "completed"}},
+            ],
+            "continuationToken": None,
+        },
+    )
+
+    result = list(tc.list_task_group_incomplete_tasks(tgid))
+    assert result == ["1", "3", "4", "6", "7"]
+
+
 def test_get_ancestors(responses, root_url):
     tc.get_task_definition.cache_clear()
     tc._get_deps.cache_clear()