From fc29b1c7f0dc0771d0051c7ee9a203552d47d6ad Mon Sep 17 00:00:00 2001
From: Florin Blanaru <florin.blanaru96@gmail.com>
Date: Mon, 22 May 2023 13:28:08 +0300
Subject: [PATCH] [CI] Update the expected CI jobs list in the update_branch
 script

---
 ci/scripts/github/update_branch.py | 19 ++++++++++--
 tests/python/ci/test_ci.py         | 49 ++++++++++--------------------
 2 files changed, 33 insertions(+), 35 deletions(-)

diff --git a/ci/scripts/github/update_branch.py b/ci/scripts/github/update_branch.py
index 9f689f6bfa65..e49d9e47ab79 100755
--- a/ci/scripts/github/update_branch.py
+++ b/ci/scripts/github/update_branch.py
@@ -87,6 +87,22 @@ def commits_query(user: str, repo: str, cursor: str = None):
     """
 
 
+EXPECTED_CI_JOBS = [
+    "cross-isa-minimal/branch",
+    "gpu/branch",
+    "hexagon/branch",
+    "arm/branch",
+    "cortexm/branch",
+    "cpu/branch",
+    "docker/branch",
+    "i386/branch",
+    "lint/branch",
+    "minimal/branch",
+    "riscv/branch",
+    "wasm/branch",
+]
+
+
 def commit_passed_ci(commit: Dict[str, Any]) -> bool:
     """
     Returns true if all of a commit's statuses are SUCCESS
@@ -111,9 +127,8 @@ def commit_passed_ci(commit: Dict[str, Any]) -> bool:
 
     # Assert that specific jobs are present in the commit statuses (i.e. don't
     # approve if CI was broken and didn't schedule a job)
-    expected_jobs = {"tvm-ci/branch"}
     job_names = {name for name, status in unified_statuses}
-    for job in expected_jobs:
+    for job in EXPECTED_CI_JOBS:
         if job not in job_names:
             # Did not find expected job name
             return False
diff --git a/tests/python/ci/test_ci.py b/tests/python/ci/test_ci.py
index f83fa40c618d..ca749665c984 100644
--- a/tests/python/ci/test_ci.py
+++ b/tests/python/ci/test_ci.py
@@ -36,6 +36,8 @@
 import scripts.github
 import scripts.jenkins
 
+from scripts.github.update_branch import EXPECTED_CI_JOBS
+
 # pylint: enable=wrong-import-position,wrong-import-order
 
 
@@ -434,8 +436,12 @@ def test_cc_reviewers(
     assert f"After filtering existing reviewers, adding: {expected_reviewers}" in proc.stdout
 
 
+def generate_good_commit_status():
+    return list([{"context": context, "state": "SUCCESS"} for context in EXPECTED_CI_JOBS])
+
+
 @parameterize_named(
-    # Missing expected tvm-ci/branch test
+    # Missing expected gpu/branch test
     missing_tvm_ci_branch=dict(
         statuses=[
             {
@@ -448,48 +454,25 @@ def test_cc_reviewers(
     ),
     # Only has the right passing test
     has_expected_test=dict(
-        statuses=[
-            {
-                "context": "tvm-ci/branch",
-                "state": "SUCCESS",
-            }
-        ],
+        statuses=generate_good_commit_status(),
         expected_rc=0,
         expected_output="Found last good commit: 123: hello",
     ),
     # Check with many statuses
     many_statuses=dict(
-        statuses=[
-            {
-                "context": "tvm-ci/branch",
-                "state": "SUCCESS",
-            },
-            {
-                "context": "tvm-ci/branch2",
-                "state": "SUCCESS",
-            },
-            {
-                "context": "tvm-ci/branch3",
-                "state": "FAILED",
-            },
+        statuses=generate_good_commit_status()
+        + [
+            {"context": "gpu/branch2", "state": "SUCCESS"},
+            {"context": "gpu/branch3", "state": "FAILED"},
         ],
         expected_rc=1,
         expected_output="No good commits found in the last 1 commits",
     ),
     many_success_statuses=dict(
-        statuses=[
-            {
-                "context": "tvm-ci/branch",
-                "state": "SUCCESS",
-            },
-            {
-                "context": "tvm-ci/branch2",
-                "state": "SUCCESS",
-            },
-            {
-                "context": "tvm-ci/branch3",
-                "state": "SUCCESS",
-            },
+        statuses=generate_good_commit_status()
+        + [
+            {"context": "gpu/branch2", "state": "SUCCESS"},
+            {"context": "gpu/branch3", "state": "SUCCESS"},
         ],
         expected_rc=0,
         expected_output="Found last good commit: 123: hello",