3 files changed, 79 insertions, 0 deletions

diff --git a/src/mongo/unittest/bson_test_util.cpp b/src/mongo/unittest/bson_test_util.cpp
index 57d93d36026..1f4e52d19ff 100644
--- a/src/mongo/unittest/bson_test_util.cpp
+++ b/src/mongo/unittest/bson_test_util.cpp
@@ -55,5 +55,12 @@ GENERATE_BSON_CMP_FUNC(BSONObj, GT, SimpleBSONObjComparator::kInstance, >);
 GENERATE_BSON_CMP_FUNC(BSONObj, GTE, SimpleBSONObjComparator::kInstance, >=);
 GENERATE_BSON_CMP_FUNC(BSONObj, NE, SimpleBSONObjComparator::kInstance, !=);
 
+GENERATE_BSON_CMP_FUNC(BSONElement, EQ, SimpleBSONElementComparator::kInstance, ==);
+GENERATE_BSON_CMP_FUNC(BSONElement, LT, SimpleBSONElementComparator::kInstance, <);
+GENERATE_BSON_CMP_FUNC(BSONElement, LTE, SimpleBSONElementComparator::kInstance, <=);
+GENERATE_BSON_CMP_FUNC(BSONElement, GT, SimpleBSONElementComparator::kInstance, >);
+GENERATE_BSON_CMP_FUNC(BSONElement, GTE, SimpleBSONElementComparator::kInstance, >=);
+GENERATE_BSON_CMP_FUNC(BSONElement, NE, SimpleBSONElementComparator::kInstance, !=);
+
 }  // namespace unittest
 }  // namespace mongo
diff --git a/src/mongo/unittest/bson_test_util.h b/src/mongo/unittest/bson_test_util.h
index 9a83a3fd482..a70af117f35 100644
--- a/src/mongo/unittest/bson_test_util.h
+++ b/src/mongo/unittest/bson_test_util.h
@@ -28,6 +28,7 @@
 
 #pragma once
 
+#include "mongo/bson/simple_bsonelement_comparator.h"
 #include "mongo/bson/simple_bsonobj_comparator.h"
 #include "mongo/unittest/unittest.h"
 
@@ -47,6 +48,16 @@
 #define ASSERT_BSONOBJ_GTE(a, b) ASSERT_BSON_COMPARISON(BSONObjGTE, a, b)
 #define ASSERT_BSONOBJ_NE(a, b) ASSERT_BSON_COMPARISON(BSONObjNE, a, b)
 
+/**
+ * Use to compare two instances of type BSONElement under the default comparator in unit tests.
+ */
+#define ASSERT_BSONELT_EQ(a, b) ASSERT_BSON_COMPARISON(BSONElementEQ, a, b)
+#define ASSERT_BSONELT_LT(a, b) ASSERT_BSON_COMPARISON(BSONElementLT, a, b)
+#define ASSERT_BSONELT_LTE(a, b) ASSERT_BSON_COMPARISON(BSONElementLTE, a, b)
+#define ASSERT_BSONELT_GT(a, b) ASSERT_BSON_COMPARISON(BSONElementGT, a, b)
+#define ASSERT_BSONELT_GTE(a, b) ASSERT_BSON_COMPARISON(BSONElementGTE, a, b)
+#define ASSERT_BSONELT_NE(a, b) ASSERT_BSON_COMPARISON(BSONElementNE, a, b)
+
 namespace mongo {
 namespace unittest {
 
@@ -64,6 +75,13 @@ DECLARE_BSON_CMP_FUNC(BSONObj, LTE);
 DECLARE_BSON_CMP_FUNC(BSONObj, GT);
 DECLARE_BSON_CMP_FUNC(BSONObj, GTE);
 DECLARE_BSON_CMP_FUNC(BSONObj, NE);
+
+DECLARE_BSON_CMP_FUNC(BSONElement, EQ);
+DECLARE_BSON_CMP_FUNC(BSONElement, LT);
+DECLARE_BSON_CMP_FUNC(BSONElement, LTE);
+DECLARE_BSON_CMP_FUNC(BSONElement, GT);
+DECLARE_BSON_CMP_FUNC(BSONElement, GTE);
+DECLARE_BSON_CMP_FUNC(BSONElement, NE);
 #undef DECLARE_BSON_CMP_FUNC
 
 }  // namespace unittest
diff --git a/src/mongo/unittest/unittest_test.cpp b/src/mongo/unittest/unittest_test.cpp
index c9d8fa7a9a8..03c00a34712 100644
--- a/src/mongo/unittest/unittest_test.cpp
+++ b/src/mongo/unittest/unittest_test.cpp
@@ -189,6 +189,60 @@ TEST(UnitTestSelfTest, BSONObjGTE) {
                             << "bar"));
 }
 
+TEST(UnitTestSelfTest, BSONElementEQ) {
+    mongo::BSONObj obj1 = BSON("foo"
+                               << "bar");
+    mongo::BSONObj obj2 = BSON("foo"
+                               << "bar");
+    ASSERT_BSONELT_EQ(obj1.firstElement(), obj2.firstElement());
+}
+
+TEST(UnitTestSelfTest, BSONElementNE) {
+    mongo::BSONObj obj1 = BSON("foo"
+                               << "bar");
+    mongo::BSONObj obj2 = BSON("foo"
+                               << "baz");
+    ASSERT_BSONELT_NE(obj1.firstElement(), obj2.firstElement());
+}
+
+TEST(UnitTestSelfTest, BSONElementLT) {
+    mongo::BSONObj obj1 = BSON("foo"
+                               << "bar");
+    mongo::BSONObj obj2 = BSON("foo"
+                               << "baz");
+    ASSERT_BSONELT_LT(obj1.firstElement(), obj2.firstElement());
+}
+
+TEST(UnitTestSelfTest, BSONElementLTE) {
+    mongo::BSONObj obj1 = BSON("foo"
+                               << "bar");
+    mongo::BSONObj obj2 = BSON("foo"
+                               << "bar");
+    mongo::BSONObj obj3 = BSON("foo"
+                               << "baz");
+    ASSERT_BSONELT_LTE(obj1.firstElement(), obj2.firstElement());
+    ASSERT_BSONELT_LTE(obj1.firstElement(), obj3.firstElement());
+}
+
+TEST(UnitTestSelfTest, BSONElementGT) {
+    mongo::BSONObj obj1 = BSON("foo"
+                               << "bar");
+    mongo::BSONObj obj2 = BSON("foo"
+                               << "baz");
+    ASSERT_BSONELT_GT(obj2.firstElement(), obj1.firstElement());
+}
+
+TEST(UnitTestSelfTest, BSONElementGTE) {
+    mongo::BSONObj obj1 = BSON("foo"
+                               << "bar");
+    mongo::BSONObj obj2 = BSON("foo"
+                               << "bar");
+    mongo::BSONObj obj3 = BSON("foo"
+                               << "baz");
+    ASSERT_BSONELT_GTE(obj3.firstElement(), obj2.firstElement());
+    ASSERT_BSONELT_GTE(obj2.firstElement(), obj1.firstElement());
+}
+
 DEATH_TEST(DeathTestSelfTest, TestDeath, "Invariant failure false") {
     invariant(false);
 }