push: add ability to generate all shadow scheduler results

Fixes #126

mozci/push.py

@@ -755,13 +755,29 @@ class Push:
             name (str): The name of the shadow scheduler to query.
 
         Returns:
-            list: All task labels that would have been scheduled.
+            set: All task labels that would have been scheduled.
         """
         index = self.index + ".source.shadow-scheduler-{}".format(name)
         task = Task.create(index=index)
         labels = task.get_artifact("public/shadow-scheduler/optimized_tasks.list")
         return set(labels.splitlines())
 
+    def generate_all_shadow_scheduler_tasks(self):
+        """Generates all tasks from all of the shadow schedulers that ran on the push.
+
+        Yields:
+            tuple: Of the form (<name>, [<label>]) where the first value is the
+            name of the shadow scheduler and the second is the set of tasks it
+            would have scheduled.
+        """
+        names = [
+            label.split("shadow-scheduler-")[1]
+            for label in self.scheduled_task_labels
+            if "shadow-scheduler" in label
+        ]
+        for name in sorted(names):
+            yield name, self.get_shadow_scheduler_tasks(name)
+
     def __repr__(self):
         return f"{super(Push, self).__repr__()} rev='{self.rev}'"
 

tests/test_push.py

@@ -1,10 +1,13 @@
 # -*- coding: utf-8 -*-
 
+from itertools import count
+
 import pytest
 
 from mozci.errors import ChildPushNotFound, ParentPushNotFound, PushNotFound
 from mozci.push import Push
 from mozci.util.hgmo import HGMO
+from mozci.util.taskcluster import get_artifact_url, get_index_url
 
 
 @pytest.fixture
@@ -335,3 +338,57 @@ def test_push_child_raises(responses):
 
     with pytest.raises(ChildPushNotFound):
         push.child
+
+
+def test_generate_all_shadow_scheduler_tasks(responses):
+    rev = "a" * 40
+    shadow_schedulers = (
+        (
+            "bar",
+            ["task-1", "task-3", "task-4"],
+        ),  # names will be generated alphabetically
+        ("foo", ["task-2", "task-4"]),
+    )
+
+    push = Push(rev)
+    responses.add(
+        responses.GET,
+        get_index_url(push.index + ".taskgraph.decision"),
+        json={"taskId": 1},
+        status=200,
+    )
+
+    id = count(2)
+    responses.add(
+        responses.GET,
+        get_artifact_url(1, "public/task-graph.json"),
+        json={
+            next(id): {"label": f"source-test-shadow-scheduler-{s[0]}"}
+            for s in shadow_schedulers
+        },
+        status=200,
+    )
+
+    id = count(2)
+    for s in shadow_schedulers:
+        s_id = next(id)
+        responses.add(
+            responses.GET,
+            get_index_url(push.index + f".source.shadow-scheduler-{s[0]}"),
+            json={"taskId": s_id},
+            status=200,
+        )
+
+        responses.add(
+            responses.GET,
+            get_artifact_url(s_id, "public/shadow-scheduler/optimized_tasks.list"),
+            stream=True,
+            body="\n".join(s[1]),
+            status=200,
+        )
+
+    # retrieve the data
+    for i, (name, tasks) in enumerate(push.generate_all_shadow_scheduler_tasks()):
+        print(i, name, tasks)
+        assert name == shadow_schedulers[i][0]
+        assert tasks == set(shadow_schedulers[i][1])