1 files changed, 34 insertions, 6 deletions

diff --git a/absl/container/btree_test.cc b/absl/container/btree_test.cc
index 9b1b6436..a933386a 100644
--- a/absl/container/btree_test.cc
+++ b/absl/container/btree_test.cc
@@ -2696,8 +2696,36 @@ struct MultiKeyComp {
   bool operator()(const MultiKey a, const int b) const { return a.i1 < b; }
 };
 
-TEST(Btree, MultiKeyEqualRange) {
-  absl::btree_set<MultiKey, MultiKeyComp> set;
+// A heterogeneous, three-way comparator that has different equivalence classes
+// for different lookup types.
+struct MultiKeyThreeWayComp {
+  using is_transparent = void;
+  absl::weak_ordering operator()(const MultiKey a, const MultiKey b) const {
+    if (a.i1 < b.i1) return absl::weak_ordering::less;
+    if (a.i1 > b.i1) return absl::weak_ordering::greater;
+    if (a.i2 < b.i2) return absl::weak_ordering::less;
+    if (a.i2 > b.i2) return absl::weak_ordering::greater;
+    return absl::weak_ordering::equivalent;
+  }
+  absl::weak_ordering operator()(const int a, const MultiKey b) const {
+    if (a < b.i1) return absl::weak_ordering::less;
+    if (a > b.i1) return absl::weak_ordering::greater;
+    return absl::weak_ordering::equivalent;
+  }
+  absl::weak_ordering operator()(const MultiKey a, const int b) const {
+    if (a.i1 < b) return absl::weak_ordering::less;
+    if (a.i1 > b) return absl::weak_ordering::greater;
+    return absl::weak_ordering::equivalent;
+  }
+};
+
+template <typename Compare>
+class BtreeMultiKeyTest : public ::testing::Test {};
+using MultiKeyComps = ::testing::Types<MultiKeyComp, MultiKeyThreeWayComp>;
+TYPED_TEST_SUITE(BtreeMultiKeyTest, MultiKeyComps);
+
+TYPED_TEST(BtreeMultiKeyTest, EqualRange) {
+  absl::btree_set<MultiKey, TypeParam> set;
 
   for (int i = 0; i < 100; ++i) {
     for (int j = 0; j < 100; ++j) {
@@ -2713,15 +2741,15 @@ TEST(Btree, MultiKeyEqualRange) {
   }
 }
 
-TEST(Btree, MultiKeyErase) {
-  absl::btree_set<MultiKey, MultiKeyComp> set = {
+TYPED_TEST(BtreeMultiKeyTest, Erase) {
+  absl::btree_set<MultiKey, TypeParam> set = {
       {1, 1}, {2, 1}, {2, 2}, {3, 1}};
   EXPECT_EQ(set.erase(2), 2);
   EXPECT_THAT(set, ElementsAre(MultiKey{1, 1}, MultiKey{3, 1}));
 }
 
-TEST(Btree, MultiKeyCount) {
-  const absl::btree_set<MultiKey, MultiKeyComp> set = {
+TYPED_TEST(BtreeMultiKeyTest, Count) {
+  const absl::btree_set<MultiKey, TypeParam> set = {
       {1, 1}, {2, 1}, {2, 2}, {3, 1}};
   EXPECT_EQ(set.count(2), 2);
 }