googletest: support the bazel test sharding protocol

PiperOrigin-RevId: 717986884

Author: Googler

Cargo.lock

@@ -14,7 +14,7 @@
 
 [package]
 name = "googletest"
-version = "0.13.0"
+version = "0.13.1"
 keywords = ["unit", "matcher", "testing", "assertions"]
 categories = ["development-tools", "development-tools::testing"]
 description = "A rich assertion and matcher library inspired by GoogleTest for C++"
@@ -31,7 +31,7 @@ authors = [
 ]
 
 [dependencies]
-googletest_macro = { path = "../googletest_macro", version = "0.13.0" }
+googletest_macro = { path = "../googletest_macro", version = "0.13.1" }
 anyhow = { version = "1", optional = true }
 num-traits = "0.2.17"
 proptest = { version = "1.2.0", optional = true }

googletest/Cargo.toml

@@ -16,3 +16,4 @@
 
 pub(crate) mod description_renderer;
 pub mod test_outcome;
+pub mod test_sharding;

googletest/src/internal/mod.rs

@@ -0,0 +1,89 @@
+/// This module implements the googletest test sharding protocol.  The Google
+/// sharding protocol consists of the following environment variables:
+///
+/// * GTEST_TOTAL_SHARDS: total number of shards.
+/// * GTEST_SHARD_INDEX: number of this shard
+/// * GTEST_SHARD_STATUS_FILE: touch this file to indicate support for sharding.
+///
+/// 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::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";
+
+/// Environment variable specifyign the index of this test shard.
+const TEST_SHARD_INDEX: &str = "GTEST_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";
+
+thread_local! {
+    static SHARDING: OnceCell<Sharding> = const { OnceCell::new() };
+}
+
+struct Sharding {
+    this_shard: u64,
+    total_shards: NonZeroU64,
+}
+
+impl Default for Sharding {
+    fn default() -> Self {
+        Self { this_shard: 0, total_shards: NonZeroU64::MIN }
+    }
+}
+
+pub fn test_should_run(test_case_hash: u64) -> bool {
+    SHARDING.with(|sharding_cell| {
+        sharding_cell.get_or_init(Sharding::from_environment).test_should_run(test_case_hash)
+    })
+}
+
+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()) };
+        let total_shards: Option<NonZeroU64> = {
+            var(OsStr::new(TEST_TOTAL_SHARDS))
+                .ok()
+                .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)) {
+                    let pathbuf = PathBuf::from(name);
+                    if let Err(e) = create_status_file(&pathbuf) {
+                        eprintln!(
+                            "failed to create {} file {}: {}",
+                            TEST_SHARD_STATUS_FILE,
+                            pathbuf.display(),
+                            e
+                        );
+                    }
+                }
+
+                Sharding { this_shard, total_shards }
+            }
+            _ => Sharding::default(),
+        }
+    }
+}
+
+fn create_status_file(path: &Path) -> std::io::Result<()> {
+    if let Some(parent) = path.parent() {
+        fs::create_dir_all(parent)?;
+    }
+
+    File::create(path).map(|_| ())
+}

googletest/src/internal/test_sharding.rs

@@ -14,7 +14,7 @@
 
 [package]
 name = "googletest_macro"
-version = "0.13.0"
+version = "0.13.1"
 keywords = ["unit", "matcher", "testing", "assertions"]
 categories = ["development-tools", "development-tools::testing"]
 description = "Procedural macros for GoogleTest Rust"

googletest_macro/Cargo.toml

@@ -76,6 +76,18 @@ pub fn gtest(
     input: proc_macro::TokenStream,
 ) -> proc_macro::TokenStream {
     let ItemFn { attrs, sig, block, .. } = parse_macro_input!(input as ItemFn);
+    let test_case_hash: u64 = {
+        use std::collections::hash_map::DefaultHasher;
+        use std::hash::{Hash, Hasher};
+        let mut h = DefaultHasher::new();
+
+        // Only consider attrs and name for stability. Changing the function body should
+        // not affect the test case distribution.
+        attrs.hash(&mut h);
+        sig.ident.hash(&mut h);
+        h.finish()
+    };
+
     let (outer_return_type, trailer) = if attrs
         .iter()
         .any(|attr| attr.path().is_ident("should_panic"))
@@ -168,10 +180,14 @@ pub fn gtest(
         #(#attrs)*
         #outer_sig -> #outer_return_type {
             #maybe_closure
-            use googletest::internal::test_outcome::TestOutcome;
-            TestOutcome::init_current_test_outcome();
-            let result: #invocation_result_type = #invocation;
-            TestOutcome::close_current_test_outcome(#result)
+            if googletest::internal::test_sharding::test_should_run(#test_case_hash) {
+                use googletest::internal::test_outcome::TestOutcome;
+                TestOutcome::init_current_test_outcome();
+                let result: #invocation_result_type = #invocation;
+                TestOutcome::close_current_test_outcome(#result)
+            } else {
+                Ok(())
+            }
             #trailer
         }
     };
@@ -184,6 +200,7 @@ pub fn gtest(
             #function
         }
     };
+
     output.into()
 }
 

googletest_macro/src/lib.rs

@@ -14,7 +14,7 @@
 
 [package]
 name = "integration_tests"
-version = "0.13.0"
+version = "0.13.1"
 description = "Integration tests for GoogleTest Rust"
 repository = "https://github.com/google/googletest-rust"
 license = "Apache-2.0"
@@ -33,13 +33,19 @@ anyhow = "1"
 indoc = "2"
 rstest = "0.18"
 rustversion = "1.0.14"
+tempfile = "3.10.1"
 tokio = { version = "1.34", features = ["time", "macros", "rt"] }
 
 [[bin]]
 name = "integration_tests"
 path = "src/integration_tests.rs"
 test = false
 
+[[bin]]
+name = "always_fails"
+path = "src/always_fails.rs"
+test = false
+
 [[bin]]
 name = "assert_predicate_with_failure"
 path = "src/assert_predicate_with_failure.rs"

integration_tests/Cargo.toml

@@ -0,0 +1,25 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+fn main() {}
+
+#[cfg(test)]
+mod tests {
+    use googletest::prelude::*;
+
+    #[gtest]
+    fn always_fails() -> Result<()> {
+        verify_that!(2, eq(3))
+    }
+}

integration_tests/src/always_fails.rs

@@ -743,6 +743,41 @@ mod tests {
         verify_that!(output, contains_regex("Did you annotate the test with gtest?"))
     }
 
+    #[gtest]
+    fn always_fails_test_always_fails() -> Result<()> {
+        let status = run_external_process("always_fails").status()?;
+
+        verify_that!(status.success(), eq(false))
+    }
+
+    #[gtest]
+    fn always_fails_test_respects_sharding() -> Result<()> {
+        fn execute_test(shard: u64, total_shards: u64) -> Result<bool> {
+            let status_file = tempfile::tempdir()?.path().join("shard_status_file");
+            let gtest_total_shards = format!("{total_shards}");
+            let gtest_shard_index = format!("{shard}");
+
+            let success = run_external_process("always_fails")
+                .env("GTEST_TOTAL_SHARDS", gtest_total_shards)
+                .env("GTEST_SHARD_INDEX", gtest_shard_index)
+                .env("GTEST_SHARD_STATUS_FILE", &status_file)
+                .status()?
+                .success();
+
+            verify_that!(status_file.exists(), eq(true))?;
+            Ok(success)
+        }
+
+        // The test case should only run in one shard.
+        let results = [execute_test(0, 3)?, execute_test(1, 3)?, execute_test(2, 3)?];
+        let successes = results.iter().filter(|b| **b).count();
+        let failures = results.iter().filter(|b| !**b).count();
+
+        expect_that!(successes, eq(2));
+        expect_that!(failures, eq(1));
+        Ok(())
+    }
+
     #[gtest]
     fn verify_true_when_true_returns_ok() {
         assert!(verify_true!("test" == "test").is_ok())

integration_tests/src/integration_tests.rs

@@ -24,6 +24,7 @@ set -e
 
 INTEGRATION_TEST_BINARIES=(
   "integration_tests"
+  "always_fails"
   "assert_predicate_with_failure"
   "assertion_failure_in_subroutine"
   "assertion_failures_with_short_structured_actual_values"