Skip to content

Commit 0298657

Browse files
Marcin RadomskiGoogle-ML-Automation
authored andcommitted
[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) is the wrapper itself. 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
1 parent 2039fa8 commit 0298657

File tree

13 files changed

+219
-10
lines changed

13 files changed

+219
-10
lines changed

MODULE.bazel

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,8 @@ bazel_dep(name = "bazel_skylib", version = "1.8.1")
1010
bazel_dep(name = "boringssl", version = "0.20250818.0")
1111
bazel_dep(name = "curl", version = "8.11.0")
1212
bazel_dep(name = "google_benchmark", version = "1.8.5", repo_name = "com_google_benchmark")
13-
bazel_dep(name = "googletest", version = "1.17.0", repo_name = "com_google_googletest")
13+
bazel_dep(name = "googletest", version = "1.17.0", repo_name = "com_google_googletest_upstream")
14+
bazel_dep(name = "xla_googletest_wrapper", version = "1.0", repo_name = "com_google_googletest")
1415
bazel_dep(name = "grpc", version = "1.74.1", repo_name = "com_github_grpc_grpc")
1516
bazel_dep(name = "gutil", version = "20250502.0", repo_name = "com_google_gutil");
1617
bazel_dep(name = "jsoncpp", version = "1.9.6", repo_name = "jsoncpp_git")
@@ -74,6 +75,11 @@ archive_override(
7475
urls = ["https://github.com/google/googletest/archive/28e9d1f26771c6517c3b4be10254887673c94018.zip"],
7576
)
7677

78+
local_path_override(
79+
module_name = "xla_googletest_wrapper",
80+
path = "third_party/xla_googletest_wrapper",
81+
)
82+
7783
##############################################################
7884
# C++ dependencies
7985

opensource_only.files

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,4 +155,6 @@ xla/third_party/tensorrt/tensorrt/include/tensorrt_config.h.tpl:
155155
xla/third_party/tensorrt/tensorrt/tensorrt_config.py.tpl:
156156
xla/third_party/tensorrt/tensorrt_configure.bzl:
157157
xla/third_party/tensorrt/workspace.bzl:
158+
xla/third_party/xla_googletest_wrapper/include/gmock/gmock.h:
159+
xla/third_party/xla_googletest_wrapper/include/gtest/gtest.h:
158160
xla/third_party/zlib.BUILD:

third_party/gutil/workspace.bzl

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,4 @@ def repo():
1414
sha256 = GUTIL_SHA256,
1515
strip_prefix = "gutil-{commit}".format(commit = GUTIL_COMMIT),
1616
urls = tf_mirror_urls("https://github.com/google/gutil/archive/{commit}.tar.gz".format(commit = GUTIL_COMMIT)),
17-
repo_mapping = {
18-
"@absl-cpp": "@com_google_absl",
19-
"@google_benchmark": "@com_google_benchmark",
20-
"@googletest": "@com_google_googletest",
21-
},
2217
)
Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package(default_visibility = ["//visibility:public"])
2+
3+
cc_library(
4+
name = "gtest",
5+
# Upstream gtest is *not* marked testonly for some reason, but
6+
# gutil:status_matchers is, and non-testonly targets can't depend on
7+
# testonly ones.
8+
#
9+
# XLA doesn't use gtest in non-testonly targets though so making this
10+
# testonly should (TM) be fine.
11+
testonly = True,
12+
hdrs = [
13+
"include/gmock/gmock.h",
14+
"include/gtest/gtest.h",
15+
],
16+
includes = [
17+
"include",
18+
],
19+
deps = [
20+
"@com_google_absl//absl/status",
21+
"@com_google_absl//absl/status:statusor",
22+
"@com_google_googletest_upstream//:gtest",
23+
],
24+
)
25+
26+
cc_library(
27+
name = "gtest_main",
28+
testonly = True,
29+
deps = [
30+
":gtest",
31+
"@com_google_googletest_upstream//:gtest_main",
32+
],
33+
)
34+
35+
alias(
36+
name = "gtest_for_library",
37+
actual = ":gtest",
38+
)
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
module(name = "xla_googletest_wrapper", version = "1.0")
2+
3+
bazel_dep(name = "abseil-cpp", version = "20250814.0", repo_name = "com_google_absl")
4+
bazel_dep(name = "googletest", version = "1.17.0", repo_name = "com_google_googletest_upstream")
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
A gtest wrapper that adds ASSERT_OK, EXPECT_OK, ASSERT_OK_AND_ASSIGN to gmock.h
2+
so that the header's provided functionality matches internal gmock.
3+
4+
The repo contains a minimal set of reexports necessary to build XLA with this as
5+
a drop-in replacement for googletest.
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
# Mark the repo as a bazel repo.
Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
"""Reexports googletest_deps from upstream googletest."""
2+
3+
# protobuf loads for @com_google_googletest//:googletest_deps.bzl so we need to
4+
# provide one in the wrapper.
5+
load("@com_google_googletest_upstream//:googletest_deps.bzl", upstream_deps = "googletest_deps")
6+
7+
def googletest_deps():
8+
upstream_deps()
Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
/* Copyright 2017 The Abseil Authors & TensorFlow Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
#ifndef GOOGLETEST_WRAPPER_GMOCK_GMOCK_H_
17+
#define GOOGLETEST_WRAPPER_GMOCK_GMOCK_H_
18+
19+
// gmock/gmock.h wrapper that also provides assert macros.
20+
//
21+
// These already exist in internal version of gmock, but upstream version
22+
// doesn't have them. We use this wrapper to make dependency translation when
23+
// exporting to OSS easier.
24+
//
25+
// - We want to use standard internal header and ASSERT_OK, EXPECT_OK macros
26+
// when developing internally.
27+
// - We want the same macros to work externally, rather than having to add or
28+
// strip TF_ prefix.
29+
// - We want the OSS export to still work after the export and header
30+
// translation.
31+
// - We want to minimize the amount of patching third party projects to reduce
32+
// maintenance overhead.
33+
// - To ensure the OSS patches cleanly apply onto internal repo, we need the
34+
// header translation to be reversible, which requires 1:1 header mapping.
35+
//
36+
// To achieve this, we swap out gmock.h for this wrapper in all XLA code, which
37+
// should (TM) make ASSERT_OK/EXPECT_OK "just work" in all XLA tests.
38+
//
39+
// The only way to make this work without patching googletest and/or absl is to
40+
// make XLA *always* use this wrapper, and *never* directly depend on upstream
41+
// googletest.
42+
//
43+
// absl/status/status_matchers.h depends on gmock.h, so we can't simply add it
44+
// here. This causes either:
45+
//
46+
// - A circular dependency between this and absl - which bazel doesn't allow,
47+
// - absl dependency on the upstream gmock - which depending on the dependency
48+
// graph structure may introduce upstream gmock include path *before* one
49+
// defined in here, so we end up with *sometimes* including the wrong one and
50+
// the entire idea of drop-in replacing gmock.h goes out of the window.
51+
52+
#include_next "gmock/gmock.h"
53+
54+
#include "absl/status/status.h"
55+
#include "absl/status/statusor.h"
56+
57+
namespace xla_testing {
58+
59+
// Macros for testing the results of functions that return absl::Status or
60+
// absl::StatusOr<T> (for any type T).
61+
#define EXPECT_OK(expression) EXPECT_THAT(expression, ::xla_testing::IsOk())
62+
#define ASSERT_OK(expression) ASSERT_THAT(expression, ::xla_testing::IsOk())
63+
64+
#define ASSERT_OK_AND_ASSIGN(lhs, rexpr) \
65+
TF_ASSERT_OK_AND_ASSIGN_IMPL( \
66+
XLA_STATUS_MACROS_CONCAT_NAME(_status_or_value, __COUNTER__), \
67+
lhs, rexpr);
68+
69+
#define ASSERT_OK_AND_ASSIGN_IMPL(statusor, lhs, rexpr) \
70+
auto statusor = (rexpr); \
71+
ASSERT_OK(statusor.status()); \
72+
lhs = std::move(statusor).value()
73+
74+
#define XLA_STATUS_MACROS_CONCAT_NAME(x, y) XLA_STATUS_MACROS_CONCAT_IMPL(x, y)
75+
#define XLA_STATUS_MACROS_CONCAT_IMPL(x, y) x##y
76+
77+
78+
namespace internal {
79+
80+
inline const absl::Status& GetStatus(const absl::Status& status) {
81+
return status;
82+
}
83+
84+
template <typename T>
85+
inline const absl::Status& GetStatus(const absl::StatusOr<T>& status) {
86+
return status.status();
87+
}
88+
89+
// Monomorphic implementation of matcher IsOk() for a given type T.
90+
// T can be Status, StatusOr<>, or a reference to either of them.
91+
template <typename T>
92+
class MonoIsOkMatcherImpl : public ::testing::MatcherInterface<T> {
93+
public:
94+
void DescribeTo(std::ostream* os) const override { *os << "is OK"; }
95+
void DescribeNegationTo(std::ostream* os) const override {
96+
*os << "is not OK";
97+
}
98+
bool MatchAndExplain(T actual_value,
99+
::testing::MatchResultListener*) const override {
100+
return GetStatus(actual_value).ok();
101+
}
102+
};
103+
104+
// Implements IsOk() as a polymorphic matcher.
105+
class IsOkMatcher {
106+
public:
107+
template <typename T>
108+
/*implicit*/ operator ::testing::Matcher<T>() const { // NOLINT
109+
return ::testing::Matcher<T>(new MonoIsOkMatcherImpl<const T&>());
110+
}
111+
};
112+
113+
} // namespace internal
114+
115+
// Returns a gMock matcher that matches a Status or StatusOr<> which is OK.
116+
inline ::xla_testing::internal::IsOkMatcher IsOk() {
117+
return ::xla_testing::internal::IsOkMatcher();
118+
}
119+
120+
} // namespace xla_testing
121+
122+
#endif // GOOGLETEST_WRAPPER_GMOCK_GMOCK_H_
Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
#ifndef GOOGLETEST_WRAPPER_INCLUDE_GTEST_GTEST_H_
17+
#define GOOGLETEST_WRAPPER_INCLUDE_GTEST_GTEST_H_
18+
19+
#include_next "gtest/gtest.h"
20+
21+
#endif // GOOGLETEST_WRAPPER_INCLUDE_GTEST_GTEST_H_

0 commit comments

Comments
 (0)