Create OSS equivalent of ParseTextProtoOrDie
PiperOrigin-RevId: 815712863
diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index bf20391..14b4fd3 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -551,10 +551,9 @@
":thunk_proto_cc",
"//xla/service:buffer_assignment",
"//xla/tsl/platform:statusor",
+ "//xla/tsl/util/proto:parse_text_proto",
"//xla/tsl/util/proto:proto_matchers",
- "@com_google_absl//absl/log:check",
"@com_google_googletest//:gtest_main",
- "@local_tsl//tsl/platform:protobuf",
],
)
@@ -714,10 +713,10 @@
"//xla/service:buffer_assignment",
"//xla/tsl/platform:statusor",
"//xla/tsl/platform:test",
+ "//xla/tsl/util/proto:parse_text_proto",
"//xla/tsl/util/proto:proto_matchers",
"@com_google_absl//absl/strings:string_view",
"@com_google_googletest//:gtest_main",
- "@local_tsl//tsl/platform:protobuf",
],
)
@@ -901,6 +900,7 @@
"//xla/tsl/lib/core:status_test_util",
"//xla/tsl/platform:statusor",
"//xla/tsl/platform:test",
+ "//xla/tsl/util/proto:parse_text_proto",
"//xla/tsl/util/proto:proto_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
@@ -1053,10 +1053,9 @@
":thunk_proto_cc",
"//xla/service:buffer_assignment",
"//xla/tsl/platform:statusor",
+ "//xla/tsl/util/proto:parse_text_proto",
"//xla/tsl/util/proto:proto_matchers",
- "@com_google_absl//absl/log:check",
"@com_google_googletest//:gtest_main",
- "@local_tsl//tsl/platform:protobuf",
],
)
@@ -1835,6 +1834,7 @@
"//xla/tsl/platform:status_matchers",
"//xla/tsl/platform:statusor",
"//xla/tsl/platform:test",
+ "//xla/tsl/util/proto:parse_text_proto",
"//xla/tsl/util/proto:proto_matchers",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
@@ -2284,12 +2284,11 @@
"//xla/service:buffer_assignment",
"//xla/tsl/platform:status_matchers",
"//xla/tsl/platform:statusor",
+ "//xla/tsl/util/proto:parse_text_proto",
"//xla/tsl/util/proto:proto_matchers",
- "@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@com_google_googletest//:gtest_main",
- "@local_tsl//tsl/platform:protobuf",
],
)
diff --git a/third_party/xla/xla/backends/gpu/runtime/copy_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/copy_thunk_test.cc
index b23af73..0be4c72 100644
--- a/third_party/xla/xla/backends/gpu/runtime/copy_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/copy_thunk_test.cc
@@ -20,18 +20,18 @@
#include <gmock/gmock.h>
#include <gtest/gtest.h>
-#include "absl/log/check.h"
#include "xla/backends/gpu/runtime/thunk.h"
#include "xla/backends/gpu/runtime/thunk.pb.h"
#include "xla/service/buffer_assignment.h"
#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/util/proto/parse_text_proto.h"
#include "xla/tsl/util/proto/proto_matchers.h"
-#include "tsl/platform/protobuf.h"
namespace xla::gpu {
namespace {
using ::tsl::proto_testing::EqualsProto;
+using ::tsl::proto_testing::ParseTextProtoOrDie;
TEST(CopyThunkTest, ToProto) {
Thunk::ThunkInfo thunk_info;
@@ -60,8 +60,7 @@
}
TEST(CopyThunkTest, FromProto) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info {
profile_annotation: "profile_annotation"
@@ -72,8 +71,7 @@
destination_buffer { offset: 0 size: 256 buffer_allocation_index: 1 }
mem_size: 256
}
- )pb",
- &proto));
+ )pb");
Thunk::ThunkInfo thunk_info;
thunk_info.profile_annotation = "profile_annotation";
@@ -128,8 +126,7 @@
}
TEST(DeviceToHostCopyThunkProtoTest, FromProto) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info {
profile_annotation: "profile_annotation"
@@ -146,8 +143,7 @@
mem_size: 256
}
}
- )pb",
- &proto));
+ )pb");
Thunk::ThunkInfo thunk_info;
thunk_info.profile_annotation = "profile_annotation";
@@ -205,8 +201,7 @@
}
TEST(HostToDeviceCopyThunkProtoTest, FromProto) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info {
profile_annotation: "profile_annotation"
@@ -223,8 +218,7 @@
mem_size: 256
}
}
- )pb",
- &proto));
+ )pb");
Thunk::ThunkInfo thunk_info;
thunk_info.profile_annotation = "profile_annotation";
@@ -280,8 +274,7 @@
}
TEST(DeviceToDeviceCopyThunkProtoTest, FromProto) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info {
profile_annotation: "profile_annotation"
@@ -298,8 +291,7 @@
mem_size: 256
}
}
- )pb",
- &proto));
+ )pb");
Thunk::ThunkInfo thunk_info;
thunk_info.profile_annotation = "profile_annotation";
diff --git a/third_party/xla/xla/backends/gpu/runtime/gemm_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/gemm_thunk_test.cc
index 47f3588..d59734a 100644
--- a/third_party/xla/xla/backends/gpu/runtime/gemm_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/gemm_thunk_test.cc
@@ -27,10 +27,11 @@
#include "xla/service/buffer_assignment.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/tsl/platform/test.h"
+#include "xla/tsl/util/proto/parse_text_proto.h"
#include "xla/tsl/util/proto/proto_matchers.h"
-#include "tsl/platform/protobuf.h"
using ::tsl::proto_testing::EqualsProto;
+using ::tsl::proto_testing::ParseTextProtoOrDie;
namespace xla::gpu {
namespace {
@@ -93,11 +94,7 @@
}
)pb";
- ThunkProto original_thunk_proto;
- ASSERT_TRUE(tsl::protobuf::TextFormat::ParseFromString(
- std::string(kProtoText), // NOLINT -- openxla protobuf version requires
- // it to be a string
- &original_thunk_proto));
+ ThunkProto original_thunk_proto = ParseTextProtoOrDie<ThunkProto>(kProtoText);
std::vector<BufferAllocation> buffer_allocations;
buffer_allocations.emplace_back(/*index=*/0, /*size=*/100, /*color=*/10);
diff --git a/third_party/xla/xla/backends/gpu/runtime/kernel_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/kernel_thunk_test.cc
index e319817..589af3d 100644
--- a/third_party/xla/xla/backends/gpu/runtime/kernel_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/kernel_thunk_test.cc
@@ -56,6 +56,7 @@
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/tsl/platform/test.h"
+#include "xla/tsl/util/proto/parse_text_proto.h"
#include "xla/tsl/util/proto/proto_matchers.h"
#include "xla/xla_data.pb.h"
@@ -63,6 +64,7 @@
namespace {
using ::tsl::proto_testing::EqualsProto;
+using ::tsl::proto_testing::ParseTextProtoOrDie;
using Kind = Thunk::Kind;
TEST(KernelThunkTest, CreateWithDefaultValues) {
@@ -421,9 +423,8 @@
}
)pb";
- ThunkProto tma_kernel_thunk_proto;
- tsl::protobuf::TextFormat::ParseFromString(tma_kernel_thunk,
- &tma_kernel_thunk_proto);
+ ThunkProto tma_kernel_thunk_proto =
+ ParseTextProtoOrDie<ThunkProto>(tma_kernel_thunk);
const size_t total_byte_size =
tma_kernel_thunk_proto.kernel_thunk().args(0).size() +
diff --git a/third_party/xla/xla/backends/gpu/runtime/memset_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/memset_thunk_test.cc
index fbf512b..a539493 100644
--- a/third_party/xla/xla/backends/gpu/runtime/memset_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/memset_thunk_test.cc
@@ -20,22 +20,21 @@
#include <gmock/gmock.h>
#include <gtest/gtest.h>
-#include "absl/log/check.h"
#include "xla/backends/gpu/runtime/thunk.h"
#include "xla/backends/gpu/runtime/thunk.pb.h"
#include "xla/service/buffer_assignment.h"
#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/util/proto/parse_text_proto.h"
#include "xla/tsl/util/proto/proto_matchers.h"
-#include "tsl/platform/protobuf.h"
namespace xla::gpu {
namespace {
using ::tsl::proto_testing::EqualsProto;
+using ::tsl::proto_testing::ParseTextProtoOrDie;
TEST(MemzeroThunkTest, ProtoRoundTrip) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info {
profile_annotation: "partition_id_profile_annotation"
@@ -44,8 +43,7 @@
memzero_thunk {
dest_buffer { offset: 0 size: 4 buffer_allocation_index: 0 }
}
- )pb",
- &proto));
+ )pb");
std::vector<BufferAllocation> buffer_allocations = {
BufferAllocation(/*index=*/0, /*size=*/4, /*color=*/0)};
diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization_test.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization_test.cc
index 8da57e5..e929e64 100644
--- a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization_test.cc
@@ -21,7 +21,6 @@
#include <gmock/gmock.h>
#include <gtest/gtest.h>
-#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "xla/backends/gpu/runtime/conditional_thunk.h"
@@ -33,8 +32,8 @@
#include "xla/service/buffer_assignment.h"
#include "xla/tsl/platform/status_matchers.h"
#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/util/proto/parse_text_proto.h"
#include "xla/tsl/util/proto/proto_matchers.h"
-#include "tsl/platform/protobuf.h"
namespace xla::gpu {
namespace {
@@ -44,6 +43,7 @@
using ::testing::Property;
using ::testing::WhenDynamicCastTo;
using ::tsl::proto_testing::EqualsProto;
+using ::tsl::proto_testing::ParseTextProtoOrDie;
using Kind = Thunk::Kind;
TEST(ThunkProtoDeserializationTest, SequentialThunkChain) {
@@ -74,8 +74,7 @@
}
TEST(ThunkProtoDeserializationTest, CopyThunk) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info {
profile_annotation: "profile_annotation"
@@ -86,8 +85,7 @@
destination_buffer { offset: 0 size: 256 buffer_allocation_index: 1 }
mem_size: 256
}
- )pb",
- &proto));
+ )pb");
std::vector<BufferAllocation> buffer_allocations = {
BufferAllocation(/*index=*/0, /*size=*/1024, /*color=*/0),
@@ -102,8 +100,7 @@
}
TEST(ThunkProtoDeserializationTest, DeviceToHostCopyThunk) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info {
profile_annotation: "profile_annotation"
@@ -120,8 +117,7 @@
mem_size: 256
}
}
- )pb",
- &proto));
+ )pb");
std::vector<BufferAllocation> buffer_allocations = {
BufferAllocation(/*index=*/0, /*size=*/1024, /*color=*/0),
@@ -136,8 +132,7 @@
}
TEST(ThunkProtoDeserializationTest, HostToDeviceCopyThunk) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info {
profile_annotation: "profile_annotation"
@@ -154,8 +149,7 @@
mem_size: 256
}
}
- )pb",
- &proto));
+ )pb");
std::vector<BufferAllocation> buffer_allocations = {
BufferAllocation(/*index=*/0, /*size=*/1024, /*color=*/0),
@@ -170,8 +164,7 @@
}
TEST(ThunkProtoDeserializationTest, DeviceToDeviceCopyThunk) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info {
profile_annotation: "profile_annotation"
@@ -188,8 +181,7 @@
mem_size: 256
}
}
- )pb",
- &proto));
+ )pb");
std::vector<BufferAllocation> buffer_allocations = {
BufferAllocation(/*index=*/0, /*size=*/1024, /*color=*/0),
@@ -204,8 +196,7 @@
}
TEST(ThunkProtoDeserializationTest, WhileThunk) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info {
profile_annotation: "profile_annotation"
@@ -263,8 +254,7 @@
}
trip_count: 10
}
- )pb",
- &proto));
+ )pb");
std::vector<BufferAllocation> buffer_allocations = {
BufferAllocation(/*index=*/0, /*size=*/1024, /*color=*/0),
@@ -283,8 +273,7 @@
}
TEST(ThunkProtoDeserializationTest, ConditionalThunk) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info {
profile_annotation: "profile_annotation"
@@ -354,8 +343,7 @@
}
branch_index_is_bool: true
}
- )pb",
- &proto));
+ )pb");
std::vector<BufferAllocation> buffer_allocations = {
BufferAllocation(/*index=*/0, /*size=*/1024, /*color=*/0),
@@ -374,13 +362,11 @@
}
TEST(ThunkProtoDeserializationTest, WaitForStreamsThunk) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info { execution_stream_id: 7 }
wait_for_streams_thunk { stream_id: 1 wait_for_stream_id: 2 }
- )pb",
- &proto));
+ )pb");
TF_ASSERT_OK_AND_ASSIGN(
std::unique_ptr<Thunk> thunk,
@@ -391,8 +377,7 @@
}
TEST(ThunkProtoDeserializationTest, CudnnThunk) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info { execution_stream_id: 7 }
cudnn_thunk {
@@ -400,8 +385,7 @@
args { buffer_allocation_index: 0 }
args { buffer_allocation_index: 1 }
}
- )pb",
- &proto));
+ )pb");
std::vector<BufferAllocation> buffer_allocations = {
BufferAllocation(/*index=*/0, /*size=*/1024, /*color=*/0),
BufferAllocation(/*index=*/1, /*size=*/1024, /*color=*/0),
@@ -415,8 +399,7 @@
}
TEST(ThunkProtoDeserializationTest, CublasLtMatmulThunk) {
- ThunkProto proto;
- ASSERT_TRUE(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info { profile_annotation: "custom-call.4" }
cublas_lt_matmul_thunk {
@@ -463,8 +446,7 @@
c { size: 161600 buffer_allocation_index: 5 }
d { size: 161600 buffer_allocation_index: 5 }
}
- )pb",
- &proto));
+ )pb");
std::vector<BufferAllocation> allocations = {
BufferAllocation(/*index=*/0, /*size=*/4, /*color=*/0), // UNUSED
@@ -482,12 +464,10 @@
}
TEST(ThunkProtoDeserializationTest, EmptyThunkImplReturnsAnError) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info { execution_stream_id: 7 }
- )pb",
- &proto));
+ )pb");
EXPECT_THAT(DeserializeThunkProto(proto, /*buffer_allocations=*/{}),
absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
diff --git a/third_party/xla/xla/backends/gpu/runtime/while_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/while_thunk_test.cc
index 615e948..bdaa348 100644
--- a/third_party/xla/xla/backends/gpu/runtime/while_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/while_thunk_test.cc
@@ -45,6 +45,7 @@
#include "xla/tsl/platform/status_matchers.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/tsl/platform/test.h"
+#include "xla/tsl/util/proto/parse_text_proto.h"
#include "xla/tsl/util/proto/proto_matchers.h"
#include "tsl/platform/protobuf.h"
@@ -53,6 +54,7 @@
using ::testing::ElementsAre;
using ::tsl::proto_testing::EqualsProto;
+using ::tsl::proto_testing::ParseTextProtoOrDie;
using ::tsl::testing::IsOk;
using Kind = Thunk::Kind;
@@ -313,8 +315,7 @@
}
TEST(WhileThunkTest, FromProto) {
- ThunkProto proto;
- CHECK(tsl::protobuf::TextFormat::ParseFromString(
+ ThunkProto proto = ParseTextProtoOrDie<ThunkProto>(
R"pb(
thunk_info {
profile_annotation: "profile_annotation"
@@ -356,8 +357,7 @@
}
trip_count: 10
}
- )pb",
- &proto));
+ )pb");
Thunk::ThunkInfo thunk_info;
thunk_info.profile_annotation = "profile_annotation";
diff --git a/third_party/xla/xla/tsl/util/proto/BUILD b/third_party/xla/xla/tsl/util/proto/BUILD
index b4f25ad..ce6a593 100644
--- a/third_party/xla/xla/tsl/util/proto/BUILD
+++ b/third_party/xla/xla/tsl/util/proto/BUILD
@@ -59,3 +59,24 @@
"@local_tsl//tsl/platform:protobuf",
],
)
+
+cc_library(
+ name = "parse_text_proto",
+ testonly = True,
+ hdrs = ["parse_text_proto.h"],
+ deps = [
+ "@com_google_absl//absl/log:check",
+ "@com_google_absl//absl/strings:string_view",
+ "@com_google_protobuf//:protobuf",
+ ],
+)
+
+tsl_cc_test(
+ name = "parse_text_proto_test",
+ srcs = ["parse_text_proto_test.cc"],
+ deps = [
+ ":parse_text_proto",
+ ":proto_matchers_test_protos_cc",
+ "@com_google_googletest//:gtest_main",
+ ],
+)
diff --git a/third_party/xla/xla/tsl/util/proto/parse_text_proto.h b/third_party/xla/xla/tsl/util/proto/parse_text_proto.h
new file mode 100644
index 0000000..c111d34
--- /dev/null
+++ b/third_party/xla/xla/tsl/util/proto/parse_text_proto.h
@@ -0,0 +1,46 @@
+/* Copyright 2025 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_TSL_UTIL_PROTO_PARSE_TEXT_PROTO_H_
+#define XLA_TSL_UTIL_PROTO_PARSE_TEXT_PROTO_H_
+
+#include "absl/log/check.h"
+#include "absl/strings/string_view.h"
+#include "google/protobuf/message.h"
+#include "google/protobuf/text_format.h"
+
+namespace tsl::proto_testing {
+
+// Parses the given `text_proto` into a protobuf message of type `T`.
+//
+// `T` must be a protobuf message type.
+//
+// This is a test-only utility that is equivalent to the Google internal
+// `ParseTextProtoOrDie`, but works in OSS. Note that you must explicitly
+// specify the template argument, unlike in the internal version, where the type
+// can be inferred.
+//
+// Usage: auto proto = ParseTextProtoOrDie<MyProto>(R"pb(...)pb");
+template <typename T>
+inline T ParseTextProtoOrDie(absl::string_view text_proto) {
+ T proto;
+ CHECK(google::protobuf::TextFormat::ParseFromString(text_proto, &proto))
+ << "Failed to parse text proto";
+ return proto;
+}
+
+} // namespace tsl::proto_testing
+
+#endif // XLA_TSL_UTIL_PROTO_PARSE_TEXT_PROTO_H_
diff --git a/third_party/xla/xla/tsl/util/proto/parse_text_proto_test.cc b/third_party/xla/xla/tsl/util/proto/parse_text_proto_test.cc
new file mode 100644
index 0000000..bc92613
--- /dev/null
+++ b/third_party/xla/xla/tsl/util/proto/parse_text_proto_test.cc
@@ -0,0 +1,43 @@
+/* Copyright 2025 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/tsl/util/proto/parse_text_proto.h"
+
+#include <gtest/gtest.h>
+#include "xla/tsl/util/proto/proto_matchers_test_protos.pb.h"
+
+namespace tsl::proto_testing {
+namespace {
+
+TEST(ParseTextProtoOrDieTest, ParsesValidTextProto) {
+ Foo foo = ParseTextProtoOrDie<Foo>(R"pb(
+ s1: "hello" i2: 42 r3: "a" r3: "b"
+ )pb");
+ EXPECT_EQ(foo.s1(), "hello");
+ EXPECT_EQ(foo.i2(), 42);
+ ASSERT_EQ(foo.r3_size(), 2);
+ EXPECT_EQ(foo.r3(0), "a");
+ EXPECT_EQ(foo.r3(1), "b");
+}
+
+TEST(ParseTextProtoOrDieDeathTest, DiesOnInvalidTextProto) {
+ EXPECT_DEATH(ParseTextProtoOrDie<Foo>(R"pb(
+ invalid_field: "hello"
+ )pb"),
+ "Failed to parse text proto");
+}
+
+} // namespace
+} // namespace tsl::proto_testing
