From ae036f8eadb4cbc45f5e90ffc54b8c393871cda8 Mon Sep 17 00:00:00 2001
From: Nathan Sobo <nathan@zed.dev>
Date: Thu, 25 Sep 2025 13:03:56 -0600
Subject: [PATCH] Read env vars in TestScheduler::many (#38897)

This allows ITERATIONS and SEED environment variables to override the
hard coded values during testing.

cc @ConradIrwin @as-cii

Release Notes:

- N/A

---------

Co-authored-by: Conrad Irwin <conrad.irwin@gmail.com>
---
 crates/scheduler/src/test_scheduler.rs | 21 ++++++++++++++++-----
 crates/scheduler/src/tests.rs          | 10 ----------
 2 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/crates/scheduler/src/test_scheduler.rs b/crates/scheduler/src/test_scheduler.rs
index d9523640f4c39f6487a0383c3edba461260379af..d56db6b83909be8aa20bfe341e7767c2931f3959 100644
--- a/crates/scheduler/src/test_scheduler.rs
+++ b/crates/scheduler/src/test_scheduler.rs
@@ -41,14 +41,26 @@ impl TestScheduler {
     }
 
     /// Run a test multiple times with sequential seeds (0, 1, 2, ...)
-    pub fn many<R>(iterations: usize, mut f: impl AsyncFnMut(Arc<TestScheduler>) -> R) -> Vec<R> {
-        (0..iterations as u64)
+    pub fn many<R>(
+        default_iterations: usize,
+        mut f: impl AsyncFnMut(Arc<TestScheduler>) -> R,
+    ) -> Vec<R> {
+        let num_iterations = std::env::var("ITERATIONS")
+            .map(|iterations| iterations.parse().unwrap())
+            .unwrap_or(default_iterations);
+
+        let seed = std::env::var("SEED")
+            .map(|seed| seed.parse().unwrap())
+            .unwrap_or(0);
+
+        (seed..num_iterations as u64)
             .map(|seed| {
                 let mut unwind_safe_f = AssertUnwindSafe(&mut f);
+                eprintln!("Running seed: {seed}");
                 match panic::catch_unwind(move || Self::with_seed(seed, &mut *unwind_safe_f)) {
                     Ok(result) => result,
                     Err(error) => {
-                        eprintln!("Failing Seed: {seed}");
+                        eprintln!("\x1b[31mFailing Seed: {seed}\x1b[0m");
                         panic::resume_unwind(error);
                     }
                 }
@@ -56,8 +68,7 @@ impl TestScheduler {
             .collect()
     }
 
-    /// Run a test once with a specific seed
-    pub fn with_seed<R>(seed: u64, f: impl AsyncFnOnce(Arc<TestScheduler>) -> R) -> R {
+    fn with_seed<R>(seed: u64, f: impl AsyncFnOnce(Arc<TestScheduler>) -> R) -> R {
         let scheduler = Arc::new(TestScheduler::new(TestSchedulerConfig::with_seed(seed)));
         let future = f(scheduler.clone());
         let result = scheduler.foreground().block_on(future);
diff --git a/crates/scheduler/src/tests.rs b/crates/scheduler/src/tests.rs
index 54ce84e8841c738316232d9e8fec16f62810385e..14f436c502a2d905be9fa9d7989073398506b506 100644
--- a/crates/scheduler/src/tests.rs
+++ b/crates/scheduler/src/tests.rs
@@ -288,16 +288,6 @@ fn test_helper_methods() {
         background.spawn(async { 10 }).await
     });
     assert_eq!(results, vec![10, 10, 10]);
-
-    // Test the with_seed method
-    let result = TestScheduler::with_seed(123, async |scheduler: Arc<TestScheduler>| {
-        let background = scheduler.background();
-
-        // Spawn a background task and wait for its result
-        let task = background.spawn(async { 99 });
-        task.await
-    });
-    assert_eq!(result, 99);
 }
 
 #[test]