Migrate to Checks API for Integration Tests

.github/workflows/start-integration-tests.yml

@@ -1,28 +1,55 @@
 name: start-integration-tests
 
 on:
-  #schedule:
-  #  - cron: '*/10 * * * *'
   workflow_dispatch:
+  merge_group:
 
 jobs:
-  # Trigger for pull requests.
-  #
+  # Auto-approve integration tests for merge queue.
+  # The tests already passed on the PR, so we don't need to run them again.
+  auto-approve:
+    if: github.event_name == 'merge_group'
+
+    runs-on:
+      group: databricks-deco-testing-runner-group
+      labels: ubuntu-latest-deco
+
+    permissions:
+      checks: write
+      contents: read
+
+    steps:
+      - name: Auto-approve Check for Merge Queue
+        uses: actions/github-script@v7
+        with:
+          script: |
+            await github.rest.checks.create({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              name: 'Integration Tests',
+              head_sha: context.sha,
+              status: 'completed',
+              conclusion: 'success',
+              output: {
+                title: 'Integration Tests',
+                summary: 'Auto-approved for merge queue (tests already passed on PR)'
+              }
+            });
+
+  # Trigger integration tests for PRs.
   # This workflow triggers the integration test workflow in a different repository.
   # It requires secrets from the "test-trigger-is" environment, which are only available to authorized users.
   trigger:
+    if: github.event_name == 'workflow_dispatch'
+
     runs-on:
       group: databricks-deco-testing-runner-group
       labels: ubuntu-latest-deco
 
     environment: "test-trigger-is"
 
-    # Only run this job for PRs from branches on the main repository and not from forks.
-    # Workflows triggered by PRs from forks don't have access to the "test-trigger-is" environment.
-    if: "${{ !github.event.pull_request.head.repo.fork }}"
-
     steps:
-      - name: Generate GitHub App Token
+      - name: Generate GitHub App Token for Workflow Trigger
         id: generate-token
         uses: actions/create-github-app-token@df432ceedc7162793a195dd1713ff69aefc7379e # v2.0.6
         with:
@@ -31,11 +58,20 @@ jobs:
           owner: ${{ secrets.ORG_NAME }}
           repositories: ${{secrets.REPO_NAME}}
 
+      - name: Generate GitHub App Token for Check Updates
+        id: generate-check-token
+        uses: actions/create-github-app-token@df432ceedc7162793a195dd1713ff69aefc7379e # v2.0.6
+        with:
+          app-id: ${{ secrets.DECO_TEST_APPROVAL_APP_ID }}
+          private-key: ${{ secrets.DECO_TEST_APPROVAL_PRIVATE_KEY }}
+          owner: databricks
+
       - name: Fetch start_integration_tests.py
         run: wget https://raw.githubusercontent.com/databricks/cli/refs/heads/main/tools/start_integration_tests.py
 
       - name: Run start_integration_tests.py
         env:
           GH_TOKEN: ${{ steps.generate-token.outputs.token }}
+          GH_CHECK_TOKEN: ${{ steps.generate-check-token.outputs.token }}
         run: |-
           python3 ./start_integration_tests.py -R ${{ secrets.ORG_NAME }}/${{secrets.REPO_NAME}} --yes

tools/start_integration_tests.py

@@ -8,10 +8,9 @@
 
 import argparse
 import json
+import os
 import subprocess
 import sys
-from pathlib import Path
-import re
 
 
 CLI_REPO = "databricks/cli"
@@ -20,14 +19,17 @@
 ALLOWED_HEAD_OWNER = {"id": "MDEyOk9yZ2FuaXphdGlvbjQ5OTgwNTI=", "login": "databricks"}
 
 
-def run(cmd):
+def run(cmd, env=None):
     sys.stderr.write("+ " + " ".join(cmd) + "\n")
-    return subprocess.run(cmd, check=True)
+    return subprocess.run(cmd, check=True, env=env)
 
 
-def run_json(cmd):
+def run_json(cmd, env=None):
     sys.stderr.write("+ " + " ".join(cmd) + "\n")
-    result = subprocess.run(cmd, stdout=subprocess.PIPE, encoding="utf-8", check=True)
+    run_env = os.environ.copy()
+    if env:
+        run_env.update(env)
+    result = subprocess.run(cmd, stdout=subprocess.PIPE, encoding="utf-8", check=True, env=run_env)
 
     try:
         return json.loads(result.stdout)
@@ -36,6 +38,38 @@ def run_json(cmd):
         raise
 
 
+def create_check(commit_sha):
+    """Create a check run for the given commit and return the check_run_id."""
+    check_token = os.environ.get("GH_CHECK_TOKEN")
+    if not check_token:
+        print("Warning: GH_CHECK_TOKEN not set, skipping check creation")
+        return None
+
+    response = run_json(
+        [
+            "gh",
+            "api",
+            "-X",
+            "POST",
+            f"/repos/{CLI_REPO}/check-runs",
+            "-f",
+            "name=Integration Tests",
+            "-f",
+            f"head_sha={commit_sha}",
+            "-f",
+            "status=queued",
+            "-f",
+            "output[title]=Integration Tests",
+            "-f",
+            "output[summary]=Tests queued and will be triggered shortly...",
+        ],
+        env={"GH_TOKEN": check_token},
+    )
+    check_run_id = response.get("id")
+    print(f"Created check run: {check_run_id}")
+    return check_run_id
+
+
 def get_approved_prs_by_non_team():
     prs = run_json(
         [
@@ -108,30 +142,42 @@ def start_job(pr_number, commit_sha, author, approved_by, workflow, repo, force=
         response = input("Start integration tests? (y/n): ")
 
     if response.lower() == "y":
-        result = run(
-            [
-                "gh",
-                "workflow",
-                "run",
-                workflow,
-                "-R",
-                repo,
-                "-F",
-                f"pull_request_number={pr_number}",
-                "-F",
-                f"commit_sha={commit_sha}",
-            ],
-        )
+        check_run_id = create_check(commit_sha)
+
+        cmd = [
+            "gh",
+            "workflow",
+            "run",
+            workflow,
+            "-R",
+            repo,
+            "--ref",
+            "main",
+            "-F",
+            f"pull_request_number={pr_number}",
+            "-F",
+            f"commit_sha={commit_sha}",
+        ]
+        if check_run_id:
+            cmd.extend(["-F", f"check_run_id={check_run_id}"])
+
+        run(cmd)
         print(f"Started integration tests for PR #{pr_number}")
 
 
 def get_status(commit_sha):
-    statuses = run_json(["gh", "api", f"repos/{CLI_REPO}/commits/{commit_sha}/statuses"])
+    response = run_json(["gh", "api", f"repos/{CLI_REPO}/commits/{commit_sha}/check-runs"])
     result = []
-    for st in statuses:
-        if st["context"] != "Integration Tests Check":
+    for check in response.get("check_runs", []):
+        if check["name"] != "Integration Tests":
             continue
-        result.append(f"{st['state']} {st['target_url']}")
+        status = check["status"]
+        conclusion = check.get("conclusion", "")
+        details_url = check.get("details_url", "")
+        if conclusion:
+            result.append(f"{conclusion} {details_url}")
+        else:
+            result.append(f"{status} {details_url}")
     return result