[XLA] Add googletest wrapper that adds status asserts to gmock.h

And replace @com_google_googletest for internal uses. The only
exceptions that still need to rely on upstream googletest (renamed to
@com_google_googletest_upstream) are the wrapper and its direct
dependencies (right now just gutil, and through it absl).

This makes it possible to use the same assert macros internally and in
OSS.

Implement TF_{ASSERT,EXPECT}_OK in terms of unprefixed ones for
verification.

PiperOrigin-RevId: 837550145

Author: Marcin Radomski

MODULE.bazel

@@ -10,7 +10,8 @@ bazel_dep(name = "bazel_skylib", version = "1.8.1")
 bazel_dep(name = "boringssl", version = "0.20250818.0")
 bazel_dep(name = "curl", version = "8.11.0")
 bazel_dep(name = "google_benchmark", version = "1.8.5", repo_name = "com_google_benchmark")
-bazel_dep(name = "googletest", version = "1.17.0", repo_name = "com_google_googletest")
+bazel_dep(name = "googletest", version = "1.17.0", repo_name = "com_google_googletest_upstream")
+bazel_dep(name = "xla_googletest_wrapper", version = "1.17.0", repo_name = "com_google_googletest")
 bazel_dep(name = "grpc", version = "1.74.1", repo_name = "com_github_grpc_grpc")
 bazel_dep(name = "gutil", version = "20250502.0", repo_name = "com_google_gutil");
 bazel_dep(name = "jsoncpp", version = "1.9.6", repo_name = "jsoncpp_git")

third_party/absl/workspace.bzl

@@ -26,6 +26,6 @@ def repo():
         ],
         repo_mapping = {
             "@google_benchmark": "@com_google_benchmark",
-            "@googletest": "@com_google_googletest",
+            "@googletest": "@com_google_googletest_upstream",
         },
     )

third_party/gutil/workspace.bzl

@@ -15,8 +15,6 @@ def repo():
         strip_prefix = "gutil-{commit}".format(commit = GUTIL_COMMIT),
         urls = tf_mirror_urls("https://github.com/google/gutil/archive/{commit}.tar.gz".format(commit = GUTIL_COMMIT)),
         repo_mapping = {
-            "@absl-cpp": "@com_google_absl",
-            "@google_benchmark": "@com_google_benchmark",
-            "@googletest": "@com_google_googletest",
+            "@com_google_googletest": "@com_google_googletest_upstream",
         },
     )

workspace2.bzl

@@ -7,7 +7,7 @@ load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe")
 load("@io_bazel_rules_closure//closure:defs.bzl", "filegroup_external")
 load("@rules_ml_toolchain//gpu/sycl:sycl_configure.bzl", "sycl_configure")
 load("@rules_ml_toolchain//gpu/sycl:sycl_init_repository.bzl", "sycl_init_repository")
-load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
+load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls", "tf_vendored")
 load("//third_party/absl:workspace.bzl", absl = "repo")
 load("//third_party/benchmark:workspace.bzl", benchmark = "repo")
 load("//third_party/clang_toolchain:cc_configure_clang.bzl", "cc_download_clang_toolchain")
@@ -305,8 +305,13 @@ def _tf_repositories():
         },
     )
 
-    tf_http_archive(
+    tf_vendored(
         name = "com_google_googletest",
+        path = "xla_googletest_wrapper",
+    )
+
+    tf_http_archive(
+        name = "com_google_googletest_upstream",
         # Use the commit on 2025/6/09:
         # https://github.com/google/googletest/commit/28e9d1f26771c6517c3b4be10254887673c94018
         sha256 = "f253ca1a07262f8efde8328e4b2c68979e40ddfcfc001f70d1d5f612c7de2974",

xla/tsl/lib/core/BUILD

@@ -43,6 +43,7 @@ cc_library(
         "//xla/tsl/platform:status_matchers",
         "//xla/tsl/platform:test",
         "@com_google_absl//absl/status:status_matchers",
+        "@com_google_googletest//:gtest_for_library",
     ],
 )
 

xla/tsl/lib/core/status_test_util.h

@@ -16,13 +16,14 @@ limitations under the License.
 #ifndef XLA_TSL_LIB_CORE_STATUS_TEST_UTIL_H_
 #define XLA_TSL_LIB_CORE_STATUS_TEST_UTIL_H_
 
+#include <gmock/gmock.h>                // IWYU pragma: keep
 #include "absl/status/status_matchers.h"  // IWYU pragma: keep
 #include "xla/tsl/platform/status_matchers.h"
 #include "xla/tsl/platform/test.h"
 
 // Macros for testing the results of functions that return tensorflow::Status.
-#define TF_EXPECT_OK(statement) EXPECT_THAT((statement), ::absl_testing::IsOk())
-#define TF_ASSERT_OK(statement) ASSERT_THAT((statement), ::absl_testing::IsOk())
+#define TF_EXPECT_OK(statement) EXPECT_OK((statement))
+#define TF_ASSERT_OK(statement) ASSERT_OK((statement))
 
 // There are no EXPECT_NOT_OK/ASSERT_NOT_OK macros since they would not
 // provide much value (when they fail, they would just print the OK status

xla_googletest_wrapper/BUILD.bazel

@@ -0,0 +1,35 @@
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "gtest",
+    # Upstream gtest is *not* marked testonly for some reason, but
+    # gutil:status_matchers is, and non-testonly targets can't depend on
+    # testonly ones.
+    #
+    # XLA doesn't use gtest in non-testonly targets though so making this
+    # testonly should (TM) be fine.
+    testonly = True,
+    hdrs = [
+        "include/gmock/gmock.h",
+        "include/gtest/gtest.h",
+    ],
+    includes = [
+        "include",
+    ],
+    deps = [
+        "@com_google_googletest_upstream//:gtest",
+        "@com_google_gutil//gutil:status_matchers",
+    ],
+)
+
+cc_library(
+    name = "gtest_main",
+    deps = [
+        "@com_google_googletest_upstream//:gtest_main",
+    ],
+)
+
+alias(
+    name = "gtest_for_library",
+    actual = ":gtest",
+)

xla_googletest_wrapper/MODULE.bazel

@@ -0,0 +1,2 @@
+# Version should match the wrapped googletest.
+module(name = "xla_googletest_wrapper", version = "1.17.0")

xla_googletest_wrapper/README.md

@@ -0,0 +1,5 @@
+A gtest wrapper that adds ASSERT_OK, EXPECT_OK, ASSERT_OK_AND_ASSIGN to gmock.h
+so that the header's provided functionality matches internal gmock.
+
+The repo contains a minimal set of reexports necessary to build XLA with this as
+a drop-in replacement for googletest.

xla_googletest_wrapper/REPO.bazel

@@ -0,0 +1 @@
+# Mark the repo as a bazel repo.