test_sharding: support alternative sharding environment variables

PiperOrigin-RevId: 718988235

Author: Googler

googletest/src/internal/test_sharding.rs

@@ -8,20 +8,20 @@
 /// See also <https://google.github.io/googletest/advanced.html>
 use std::cell::OnceCell;
 use std::env::{var, var_os};
-use std::ffi::OsStr;
+use std::ffi::{OsStr, OsString};
 use std::fs::{self, File};
 use std::num::NonZeroU64;
 use std::path::{Path, PathBuf};
 
 /// Environment variable specifying the total number of test shards.
-const TEST_TOTAL_SHARDS: &str = "GTEST_TOTAL_SHARDS";
+const TEST_TOTAL_SHARDS: &[&str] = &["GTEST_TOTAL_SHARDS", "TEST_TOTAL_SHARDS"];
 
 /// Environment variable specifyign the index of this test shard.
-const TEST_SHARD_INDEX: &str = "GTEST_SHARD_INDEX";
+const TEST_SHARD_INDEX: &[&str] = &["GTEST_SHARD_INDEX", "TEST_SHARD_INDEX"];
 
 /// Environment variable specifying the name of the file we create (or cause a
 /// timestamp change on) to indicate that we support the sharding protocol.
-const TEST_SHARD_STATUS_FILE: &str = "GTEST_SHARD_STATUS_FILE";
+const TEST_SHARD_STATUS_FILE: &[&str] = &["GTEST_SHARD_STATUS_FILE", "TEST_SHARD_STATUS_FILE"];
 
 thread_local! {
     static SHARDING: OnceCell<Sharding> = const { OnceCell::new() };
@@ -44,29 +44,47 @@ pub fn test_should_run(test_case_hash: u64) -> bool {
     })
 }
 
+fn get_var(keys: &[&str]) -> Option<String> {
+    for key in keys {
+        if let Ok(value) = var(OsStr::new(key)) {
+            return Some(value);
+        }
+    }
+
+    None
+}
+
+fn get_var_os(keys: &[&str]) -> Option<OsString> {
+    for key in keys {
+        if let Some(value) = var_os(OsStr::new(key)) {
+            return Some(value);
+        }
+    }
+
+    None
+}
+
 impl Sharding {
     fn test_should_run(&self, test_case_hash: u64) -> bool {
         (test_case_hash % self.total_shards.get()) == self.this_shard
     }
 
     fn from_environment() -> Sharding {
         let this_shard: Option<u64> =
-            { var(OsStr::new(TEST_SHARD_INDEX)).ok().and_then(|value| value.parse().ok()) };
+            { get_var(TEST_SHARD_INDEX).and_then(|value| value.parse().ok()) };
         let total_shards: Option<NonZeroU64> = {
-            var(OsStr::new(TEST_TOTAL_SHARDS))
-                .ok()
+            get_var(TEST_TOTAL_SHARDS)
                 .and_then(|value| value.parse().ok())
                 .and_then(NonZeroU64::new)
         };
 
         match (this_shard, total_shards) {
             (Some(this_shard), Some(total_shards)) if this_shard < total_shards.get() => {
-                if let Some(name) = var_os(OsStr::new(TEST_SHARD_STATUS_FILE)) {
+                if let Some(name) = get_var_os(TEST_SHARD_STATUS_FILE) {
                     let pathbuf = PathBuf::from(name);
                     if let Err(e) = create_status_file(&pathbuf) {
                         eprintln!(
-                            "failed to create {} file {}: {}",
-                            TEST_SHARD_STATUS_FILE,
+                            "failed to create $GTEST_SHARD_STATUS_FILE file {}: {}",
                             pathbuf.display(),
                             e
                         );