From c13e9c1eba8f4932789a00ac1f8a9facaedc145d Mon Sep 17 00:00:00 2001
From: Brandon Williams <brandon@mystenlabs.com>
Date: Tue, 4 Jun 2024 12:16:32 -0500
Subject: [PATCH] [proptest] config: properly read environment

---
 proptest/src/test_runner/config.rs | 96 ++++++++++++++++++++----------
 1 file changed, 64 insertions(+), 32 deletions(-)

diff --git a/proptest/src/test_runner/config.rs b/proptest/src/test_runner/config.rs
index f6c98cc6..65ee1628 100644
--- a/proptest/src/test_runner/config.rs
+++ b/proptest/src/test_runner/config.rs
@@ -55,6 +55,27 @@ const DISABLE_FAILURE_PERSISTENCE: &str =
 /// Without the `std` feature this function returns config unchanged.
 #[cfg(feature = "std")]
 pub fn contextualize_config(mut result: Config) -> Config {
+    use std::env;
+    use std::ffi::OsString;
+    use std::fmt;
+    use std::str::FromStr;
+
+    const CASES: &str = "PROPTEST_CASES";
+    const MAX_LOCAL_REJECTS: &str = "PROPTEST_MAX_LOCAL_REJECTS";
+    const MAX_GLOBAL_REJECTS: &str = "PROPTEST_MAX_GLOBAL_REJECTS";
+    const MAX_FLAT_MAP_REGENS: &str = "PROPTEST_MAX_FLAT_MAP_REGENS";
+    const MAX_SHRINK_TIME: &str = "PROPTEST_MAX_SHRINK_TIME";
+    const MAX_SHRINK_ITERS: &str = "PROPTEST_MAX_SHRINK_ITERS";
+    const MAX_DEFAULT_SIZE_RANGE: &str = "PROPTEST_MAX_DEFAULT_SIZE_RANGE";
+    #[cfg(feature = "fork")]
+    const FORK: &str = "PROPTEST_FORK";
+    #[cfg(feature = "timeout")]
+    const TIMEOUT: &str = "PROPTEST_TIMEOUT";
+    const VERBOSE: &str = "PROPTEST_VERBOSE";
+    const RNG_ALGORITHM: &str = "PROPTEST_RNG_ALGORITHM";
+    const DISABLE_FAILURE_PERSISTENCE: &str =
+        "PROPTEST_DISABLE_FAILURE_PERSISTENCE";
+
     fn parse_or_warn<T: FromStr + fmt::Display>(
         src: &OsString,
         dst: &mut T,
@@ -83,66 +104,77 @@ pub fn contextualize_config(mut result: Config) -> Config {
     for (var, value) in
         env::vars_os().filter_map(|(k, v)| k.into_string().ok().map(|k| (k, v)))
     {
-        match var.as_str() {
-            CASES => parse_or_warn(&value, &mut result.cases, "u32", CASES),
-            MAX_LOCAL_REJECTS => parse_or_warn(
+        let var = var.as_str();
+
+        #[cfg(feature = "fork")]
+        if var == FORK {
+            parse_or_warn(&value, &mut result.fork, "bool", FORK);
+            continue;
+        }
+
+        #[cfg(feature = "timeout")]
+        if var == TIMEOUT {
+            parse_or_warn(&value, &mut result.timeout, "timeout", TIMEOUT);
+            continue;
+        }
+
+        if var == CASES {
+            parse_or_warn(&value, &mut result.cases, "u32", CASES);
+        } else if var == MAX_LOCAL_REJECTS {
+            parse_or_warn(
                 &value,
                 &mut result.max_local_rejects,
                 "u32",
                 MAX_LOCAL_REJECTS,
-            ),
-            MAX_GLOBAL_REJECTS => parse_or_warn(
+            );
+        } else if var == MAX_GLOBAL_REJECTS {
+            parse_or_warn(
                 &value,
                 &mut result.max_global_rejects,
                 "u32",
                 MAX_GLOBAL_REJECTS,
-            ),
-            MAX_FLAT_MAP_REGENS => parse_or_warn(
+            );
+        } else if var == MAX_FLAT_MAP_REGENS {
+            parse_or_warn(
                 &value,
                 &mut result.max_flat_map_regens,
                 "u32",
                 MAX_FLAT_MAP_REGENS,
-            ),
-            #[cfg(feature = "fork")]
-            FORK => parse_or_warn(&value, &mut result.fork, "bool", FORK),
-            #[cfg(feature = "timeout")]
-            TIMEOUT => {
-                parse_or_warn(&value, &mut result.timeout, "timeout", TIMEOUT)
-            }
-            MAX_SHRINK_TIME => parse_or_warn(
+            );
+        } else if var == MAX_SHRINK_TIME {
+            parse_or_warn(
                 &value,
                 &mut result.max_shrink_time,
                 "u32",
                 MAX_SHRINK_TIME,
-            ),
-            MAX_SHRINK_ITERS => parse_or_warn(
+            );
+        } else if var == MAX_SHRINK_ITERS {
+            parse_or_warn(
                 &value,
                 &mut result.max_shrink_iters,
                 "u32",
                 MAX_SHRINK_ITERS,
-            ),
-            MAX_DEFAULT_SIZE_RANGE => parse_or_warn(
+            );
+        } else if var == MAX_DEFAULT_SIZE_RANGE {
+            parse_or_warn(
                 &value,
                 &mut result.max_default_size_range,
                 "usize",
                 MAX_DEFAULT_SIZE_RANGE,
-            ),
-            VERBOSE => {
-                parse_or_warn(&value, &mut result.verbose, "u32", VERBOSE)
-            }
-            RNG_ALGORITHM => parse_or_warn(
+            );
+        } else if var == VERBOSE {
+            parse_or_warn(&value, &mut result.verbose, "u32", VERBOSE);
+        } else if var == RNG_ALGORITHM {
+            parse_or_warn(
                 &value,
                 &mut result.rng_algorithm,
                 "RngAlgorithm",
                 RNG_ALGORITHM,
-            ),
-            DISABLE_FAILURE_PERSISTENCE => result.failure_persistence = None,
-
-            _ => {
-                if var.starts_with("PROPTEST_") {
-                    eprintln!("proptest: Ignoring unknown env-var {}.", var);
-                }
-            }
+            );
+        } else if var == DISABLE_FAILURE_PERSISTENCE {
+            result.failure_persistence = None;
+        } else if var.starts_with("PROPTEST_") {
+            eprintln!("proptest: Ignoring unknown env-var {}.", var);
         }
     }