From cb8ad443db113b1a61f6c95714c794538f0331bb Mon Sep 17 00:00:00 2001
From: Dennis Felsing <dennis@felsing.org>
Date: Sun, 31 Dec 2023 12:39:42 +0000
Subject: [PATCH] ci: Support CI_TEST_SELECTION env variable

to directly run a selection of tests without manual selection in
buildkite. to be used from https://github.com/MaterializeInc/trigger-ci
---
 bin/ci-builder   |  3 ++-
 ci/mkpipeline.py | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/bin/ci-builder b/bin/ci-builder
index 29767d1aa952..39d442c6102c 100755
--- a/bin/ci-builder
+++ b/bin/ci-builder
@@ -171,6 +171,7 @@ case "$cmd" in
             --env AWS_SESSION_TOKEN
             --env CI
             --env CI_COVERAGE_ENABLED
+            --env CI_TEST_SELECTION
             --env CONFLUENT_CLOUD_DEVEX_KAFKA_PASSWORD
             --env CONFLUENT_CLOUD_DEVEX_KAFKA_USERNAME
             --env GITHUB_TOKEN
@@ -179,6 +180,7 @@ case "$cmd" in
             --env LAUNCHDARKLY_SDK_KEY
             --env NIGHTLY_CANARY_APP_PASSWORD
             --env MZ_CLI_APP_PASSWORD
+            --env MZ_SOFT_ASSERTIONS
             --env NO_COLOR
             --env NPM_TOKEN
             --env POLAR_SIGNALS_API_TOKEN
@@ -190,7 +192,6 @@ case "$cmd" in
             --env SCHEMA_REGISTRY_URL
             --env POSTGRES_URL
             --env COCKROACH_URL
-            --env MZ_SOFT_ASSERTIONS
         )
 
         if [[ $detach_container == "true" ]]; then
diff --git a/ci/mkpipeline.py b/ci/mkpipeline.py
index da47bc37896b..7237195648a4 100644
--- a/ci/mkpipeline.py
+++ b/ci/mkpipeline.py
@@ -134,7 +134,10 @@ def visit(step: dict[str, Any]) -> None:
 
     permit_rerunning_successful_steps(pipeline)
 
-    add_test_selection_block(pipeline, args.pipeline)
+    if test_selection := os.getenv("CI_TEST_SELECTION"):
+        trim_test_selection(pipeline, set(test_selection.split(",")))
+    else:
+        add_test_selection_block(pipeline, args.pipeline)
 
     check_depends_on(pipeline, args.pipeline)
 
@@ -240,6 +243,32 @@ def visit(step: dict[str, Any]) -> None:
                 visit(inner_step)
 
 
+def trim_test_selection(pipeline: Any, steps_to_run: set[str]) -> None:
+    def visit(step: dict[str, Any]) -> None:
+        if (
+            step.get("id") not in steps_to_run
+            and "prompt" not in step
+            and "wait" not in step
+            and "group" not in step
+            and step.get("id")
+            not in (
+                "coverage-pr-analyze",
+                "analyze",
+                "build-x86_64",
+                "build-aarch64",
+                "build-wasm",
+            )
+            and not step.get("async")
+        ):
+            step["skip"] = True
+
+    for step in pipeline["steps"]:
+        visit(step)
+        if "group" in step:
+            for inner_step in step.get("steps", []):
+                visit(inner_step)
+
+
 def add_test_selection_block(pipeline: Any, pipeline_name: str) -> None:
     selection_step = {
         "prompt": "What tests would you like to run? As a convenience, leaving all tests unchecked will run all tests.",