Fix runner type assignment in selective checks (apache#57254)

* Fix runner type assignment in selective checks

* Update notification workflow

Author: gopidesupavan

.github/workflows/ci-notification.yml

@@ -37,7 +37,7 @@ jobs:
     strategy:
       matrix:
         branch: ["v3-1-test"]
-        workflow-id: ["ci-amd.yml", "ci-arm.yml"]
+        workflow-id: ["ci-amd-arm.yml"]
     runs-on: ubuntu-latest
     steps:
       - name: "Checkout ${{ github.ref }} ( ${{ github.sha }} )"

dev/breeze/src/airflow_breeze/global_constants.py

@@ -40,9 +40,10 @@
 PUBLIC_AMD_RUNNERS = '["ubuntu-22.04"]'
 PUBLIC_ARM_RUNNERS = '["ubuntu-22.04-arm"]'
 
-RUNNERS_TYPE_MAPPING = {
-    "ubuntu-22.04": '["ubuntu-22.04"]',
-    "ubuntu-22.04-arm": '["ubuntu-22.04-arm"]',
+# The runner type cross-mapping is intentional — if the previous scheduled build used AMD, the current scheduled build should run with ARM.
+RUNNERS_TYPE_CROSS_MAPPING = {
+    "ubuntu-22.04": '["ubuntu-22.04-arm"]',
+    "ubuntu-22.04-arm": '["ubuntu-22.04"]',
 }
 
 ANSWER = ""

dev/breeze/src/airflow_breeze/utils/selective_checks.py

@@ -48,7 +48,7 @@
     PROVIDERS_COMPATIBILITY_TESTS_MATRIX,
     PUBLIC_AMD_RUNNERS,
     PUBLIC_ARM_RUNNERS,
-    RUNNERS_TYPE_MAPPING,
+    RUNNERS_TYPE_CROSS_MAPPING,
     TESTABLE_CORE_INTEGRATIONS,
     TESTABLE_PROVIDERS_INTEGRATIONS,
     GithubEvents,
@@ -1318,7 +1318,6 @@ def get_job_label(self, event_type: str, branch: str):
         if response.status_code != 200:
             get_console().print(f"[red]Error while listing workflow runs error: {response.json()}.\n")
             return None
-        get_console().print(f"[blue]Response received for workflow run {response.json()}.\n")
         runs = response.json().get("workflow_runs", [])
         if not runs:
             get_console().print(
@@ -1344,9 +1343,8 @@ def runner_type(self):
         if self._github_event in [GithubEvents.SCHEDULE, GithubEvents.PUSH]:
             branch = self._github_context_dict.get("ref_name", "main")
             label = self.get_job_label(event_type=str(self._github_event.value), branch=branch)
-            if not label:
-                return PUBLIC_AMD_RUNNERS
-            return RUNNERS_TYPE_MAPPING[label]
+
+            return RUNNERS_TYPE_CROSS_MAPPING[label] if label else PUBLIC_AMD_RUNNERS
 
         return PUBLIC_AMD_RUNNERS
 

dev/breeze/tests/test_selective_checks.py

@@ -19,7 +19,7 @@
 import json
 import re
 from typing import Any
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import pytest
 from rich.console import Console
@@ -31,6 +31,7 @@
     DEFAULT_PYTHON_MAJOR_MINOR_VERSION,
     NUMBER_OF_LOW_DEP_SLICES,
     PROVIDERS_COMPATIBILITY_TESTS_MATRIX,
+    PUBLIC_AMD_RUNNERS,
     GithubEvents,
 )
 from airflow_breeze.utils.functools_cache import clearable_cache
@@ -2478,3 +2479,69 @@ def test_ui_english_translation_changed_allowed_with_label():
         default_branch="main",
     )
     assert selective_checks.ui_english_translation_changed is True
+
+
+@patch("requests.get")
+@patch.dict("os.environ", {"GITHUB_TOKEN": "test_token"})
+def test_get_job_label(mock_get):
+    selective_checks = SelectiveChecks(
+        files=(),
+        github_event=GithubEvents.PULL_REQUEST,
+        github_repository="apache/airflow",
+        github_context_dict={},
+    )
+
+    workflow_response = Mock()
+    workflow_response.status_code = 200
+    workflow_response.json.return_value = {"workflow_runs": [{"jobs_url": "https://api.github.com/jobs/123"}]}
+
+    jobs_response = Mock()
+    jobs_response.json.return_value = {
+        "jobs": [
+            {"name": "Basic tests (ubuntu-22.04)", "labels": ["ubuntu-22.04"]},
+            {"name": "Other job", "labels": ["ubuntu-22.04"]},
+        ]
+    }
+
+    mock_get.side_effect = [workflow_response, jobs_response]
+
+    result = selective_checks.get_job_label("push", "main")
+
+    assert result == "ubuntu-22.04"
+
+
+def test_runner_type_pr():
+    selective_checks = SelectiveChecks(github_event=GithubEvents.PULL_REQUEST)
+
+    result = selective_checks.runner_type
+
+    assert result == PUBLIC_AMD_RUNNERS
+
+
+@patch("requests.get")
+@patch.dict("os.environ", {"GITHUB_TOKEN": "test_token"})
+def test_runner_type_schedule(mock_get):
+    selective_checks = SelectiveChecks(
+        files=(),
+        github_event=GithubEvents.SCHEDULE,
+        github_repository="apache/airflow",
+        github_context_dict={},
+    )
+
+    workflow_response = Mock()
+    workflow_response.status_code = 200
+    workflow_response.json.return_value = {"workflow_runs": [{"jobs_url": "https://api.github.com/jobs/123"}]}
+
+    jobs_response = Mock()
+    jobs_response.json.return_value = {
+        "jobs": [
+            {"name": "Basic tests (ubuntu-22.04)", "labels": ["ubuntu-22.04"]},
+            {"name": "Other job", "labels": ["ubuntu-22.04"]},
+        ]
+    }
+
+    mock_get.side_effect = [workflow_response, jobs_response]
+
+    result = selective_checks.runner_type
+
+    assert result == '["ubuntu-22.04-arm"]'