Read env vars in TestScheduler::many (#38897)

Nathan Sobo and Conrad Irwin created 1 month ago

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 <[e-mail]>

Change summary

crates/scheduler/src/test_scheduler.rs | 21 ++++++++++++++++-----
crates/scheduler/src/tests.rs          | 10 ----------
2 files changed, 16 insertions(+), 15 deletions(-)

Detailed changes

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);

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]