SERVER-21941: Improve handling of timestamps in assert comparison functions

2 files changed, 75 insertions, 4 deletions

diff --git a/jstests/noPassthrough/shell_assertions.js b/jstests/noPassthrough/shell_assertions.js
index 75a0a7ee47e..8915d087a3e 100644
--- a/jstests/noPassthrough/shell_assertions.js
+++ b/jstests/noPassthrough/shell_assertions.js
@@ -430,6 +430,24 @@
         });
     });
 
+    tests.push(function ltPassesWhenLessThanWithTimestamps() {
+        assert.doesNotThrow(() => {
+            assert.lt(Timestamp(3, 0), Timestamp(10, 0));
+        });
+    });
+
+    tests.push(function ltFailsWhenNotLessThanWithTimestamps() {
+        assert.throws(() => {
+            assert.lt(Timestamp(0, 10), Timestamp(0, 3));
+        });
+    });
+
+    tests.push(function ltFailsWhenEqualWithTimestamps() {
+        assert.throws(() => {
+            assert.lt(Timestamp(5, 0), Timestamp(5, 0));
+        });
+    });
+
     /* assert.gt tests */
 
     tests.push(function gtPassesWhenGreaterThan() {
@@ -490,6 +508,24 @@
         });
     });
 
+    tests.push(function gtePassesWhenGreaterThanWithTimestamps() {
+        assert.doesNotThrow(() => {
+            assert.gte(Timestamp(0, 10), Timestamp(0, 3));
+        });
+    });
+
+    tests.push(function gteFailsWhenNotGreaterThanWithTimestamps() {
+        assert.throws(() => {
+            assert.gte(Timestamp(0, 3), Timestamp(0, 10));
+        });
+    });
+
+    tests.push(function gtePassesWhenEqualWIthTimestamps() {
+        assert.doesNotThrow(() => {
+            assert.gte(Timestamp(5, 0), Timestamp(5, 0));
+        });
+    });
+
     /* assert.betweenIn tests */
 
     tests.push(function betweenInPassWhenNumberIsBetween() {
diff --git a/src/mongo/shell/assert.js b/src/mongo/shell/assert.js
index 4f8d25d9eb6..c1fea661cb5 100644
--- a/src/mongo/shell/assert.js
+++ b/src/mongo/shell/assert.js
@@ -706,10 +706,37 @@ assert = (function() {
         doassert("supposed to be null (" + (_processMsg(msg) || "") + ") was: " + tojson(what));
     };
 
+    function _shouldUseBsonWoCompare(a, b) {
+        const bsonTypes = [
+            Timestamp,
+        ];
+
+        if (typeof a !== "object" || typeof b !== "object") {
+            return false;
+        }
+
+        for (let t of bsonTypes) {
+            if (a instanceof t && b instanceof t) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
+    function _compare(f, a, b) {
+        if (_shouldUseBsonWoCompare(a, b)) {
+            const result = bsonWoCompare({_: a}, {_: b});
+            return f(result, 0);
+        }
+
+        return f(a, b);
+    }
+
     function _assertCompare(f, a, b, description, msg) {
         _validateAssertionMessage(msg);
 
-        if (f(a, b)) {
+        if (_compare(f, a, b)) {
             return;
         }
 
@@ -743,9 +770,17 @@ assert = (function() {
     assert.between = function(a, b, c, msg, inclusive) {
         _validateAssertionMessage(msg);
 
-        if ((inclusive == undefined || inclusive == true) && a <= b && b <= c) {
-            return;
-        } else if (a < b && b < c) {
+        let compareFn = (a, b) => {
+            return a < b;
+        };
+
+        if ((inclusive == undefined || inclusive == true)) {
+            compareFn = (a, b) => {
+                return a <= b;
+            };
+        }
+
+        if (_compare(compareFn, a, b) && _compare(compareFn, b, c)) {
             return;
         }