3 files changed, 144 insertions, 4 deletions
diff --git a/jstests/aggregation/sources/merge/merge_with_dollar_fields.js b/jstests/aggregation/sources/merge/merge_with_dollar_fields.js
new file mode 100644
index 00000000000..b08383fd73c
--- /dev/null
+++ b/jstests/aggregation/sources/merge/merge_with_dollar_fields.js
@@ -0,0 +1,132 @@
+// Tests $merge over documents with $-field in it.
+//
+// Sharded collections have special requirements on the join field.
+// @tags: [assumes_unsharded_collection]
+
+(function() {
+"use strict";
+
+load("jstests/libs/collection_drop_recreate.js");  // For assertDropCollection.
+
+const sourceName = 'merge_with_dollar_fields_source';
+const source = db[sourceName];
+const targetName = 'merge_with_dollar_fields_target';
+const target = db[targetName];
+
+const joinField = 'joinField';
+const sourceDoc = {
+    $dollar: 1,
+    joinField
+};
+const targetDoc = {
+    a: 1,
+    joinField
+};
+assertDropCollection(db, sourceName);
+assert.commandWorked(source.insert(sourceDoc));
+
+function runTest({whenMatched, whenNotMatched}, targetDocs) {
+    assertDropCollection(db, targetName);
+    assert.commandWorked(target.createIndex({joinField: 1}, {unique: true}));
+    assert.commandWorked(target.insert(targetDocs));
+    source.aggregate([
+        {$project: {_id: 0}},
+        {
+            $merge: {
+                into: targetName,
+                on: joinField,
+                whenMatched,
+                whenNotMatched,
+            }
+        }
+    ]);
+    return target.findOne({}, {_id: 0});
+}
+
+function runTestMatched(mode) {
+    return runTest(mode, [targetDoc]);
+}
+
+function runTestNotMatched(mode) {
+    return runTest(mode, []);
+}
+
+// TODO: SERVER-76999: Currently $merge may throw 'FailedToParse' error due to non-local updates.
+// We should return consistent results for dollar field documents.
+
+// whenMatched: 'replace', whenNotMatched: 'insert'
+assert.throwsWithCode(() => runTestMatched({whenMatched: 'replace', whenNotMatched: 'insert'}),
+                      [ErrorCodes.DollarPrefixedFieldName, ErrorCodes.FailedToParse]);
+
+try {
+    assert.docEq(sourceDoc, runTestNotMatched({whenMatched: 'replace', whenNotMatched: 'insert'}));
+} catch (error) {
+    assert.commandFailedWithCode(error, ErrorCodes.FailedToParse);
+}
+
+// whenMatched: 'replace', whenNotMatched: 'fail'
+assert.throwsWithCode(() => runTestMatched({whenMatched: 'replace', whenNotMatched: 'fail'}),
+                      [ErrorCodes.DollarPrefixedFieldName, ErrorCodes.FailedToParse]);
+
+assert.throwsWithCode(() => runTestNotMatched({whenMatched: 'replace', whenNotMatched: 'fail'}),
+                      [ErrorCodes.MergeStageNoMatchingDocument, ErrorCodes.FailedToParse]);
+
+// whenMatched: 'replace', whenNotMatched: 'discard'
+assert.throwsWithCode(() => runTestMatched({whenMatched: 'replace', whenNotMatched: 'discard'}),
+                      [ErrorCodes.DollarPrefixedFieldName, ErrorCodes.FailedToParse]);
+
+try {
+    assert.eq(null, runTestNotMatched({whenMatched: 'replace', whenNotMatched: 'discard'}));
+} catch (error) {
+    assert.commandFailedWithCode(error, ErrorCodes.FailedToParse);
+}
+
+// whenMatched: 'merge', whenNotMatched: 'insert'
+assert.throwsWithCode(() => runTestMatched({whenMatched: 'merge', whenNotMatched: 'insert'}),
+                      ErrorCodes.DollarPrefixedFieldName);
+
+assert.docEq(sourceDoc, runTestNotMatched({whenMatched: 'merge', whenNotMatched: 'insert'}));
+
+// whenMatched: 'merge', whenNotMatched: 'fail'
+assert.throwsWithCode(() => runTestMatched({whenMatched: 'merge', whenNotMatched: 'fail'}),
+                      ErrorCodes.DollarPrefixedFieldName);
+
+assert.throwsWithCode(() => runTestNotMatched({whenMatched: 'merge', whenNotMatched: 'fail'}),
+                      ErrorCodes.MergeStageNoMatchingDocument);
+
+// whenMatched: 'merge', whenNotMatched: 'discard'
+assert.throwsWithCode(() => runTestMatched({whenMatched: 'merge', whenNotMatched: 'discard'}),
+                      ErrorCodes.DollarPrefixedFieldName);
+
+assert.eq(null, runTestNotMatched({whenMatched: 'merge', whenNotMatched: 'discard'}));
+
+// whenMatched: 'keepExisting', whenNotMatched: 'insert'
+assert.docEq(targetDoc, runTestMatched({whenMatched: 'keepExisting', whenNotMatched: 'insert'}));
+
+assert.docEq(sourceDoc, runTestNotMatched({whenMatched: 'keepExisting', whenNotMatched: 'insert'}));
+
+// whenMatched: 'fail', whenNotMatched: 'insert'
+assert.throwsWithCode(() => runTestMatched({whenMatched: 'fail', whenNotMatched: 'insert'}),
+                      ErrorCodes.DuplicateKey);
+
+assert.docEq(sourceDoc, runTestNotMatched({whenMatched: 'fail', whenNotMatched: 'insert'}));
+
+// whenMatched: 'pipeline', whenNotMatched: 'insert'
+const pipeline = [{$addFields: {b: 1}}];
+const targetDocAddFields = Object.assign({}, targetDoc, {b: 1});
+assert.docEq(targetDocAddFields, runTestMatched({whenMatched: pipeline, whenNotMatched: 'insert'}));
+
+assert.docEq(sourceDoc, runTestNotMatched({whenMatched: pipeline, whenNotMatched: 'insert'}));
+
+// whenMatched: 'pipeline', whenNotMatched: 'fail'
+assert.docEq(targetDocAddFields, runTestMatched({whenMatched: pipeline, whenNotMatched: 'fail'}));
+
+assert.throwsWithCode(() => runTestNotMatched({whenMatched: pipeline, whenNotMatched: 'fail'}),
+                      ErrorCodes.MergeStageNoMatchingDocument);
+
+// whenMatched: 'pipeline', whenNotMatched: 'discard'
+assert.docEq(targetDocAddFields,
+             runTestMatched({whenMatched: pipeline, whenNotMatched: 'discard'}));
+
+assert.eq(null, runTestNotMatched({whenMatched: pipeline, whenNotMatched: 'discard'}));
+}());
diff --git a/jstests/aggregation/sources/merge/mode_merge_fail.js b/jstests/aggregation/sources/merge/mode_merge_fail.js
index 7235c8e1c7e..fcc13f8e871 100644
--- a/jstests/aggregation/sources/merge/mode_merge_fail.js
+++ b/jstests/aggregation/sources/merge/mode_merge_fail.js
@@ -97,9 +97,13 @@ const pipeline = [mergeStage];
 // and updated.
 (function testMergeUnorderedBatchUpdate() {
     const maxBatchSize = 16 * 1024 * 1024;  // 16MB
-    const docSize = 1024 * 1024;            // 1MB
+
+    // Each document is just under 1MB in order to allow for some extra space for writes that need
+    // to be serialized over the wire in certain cluster configurations. Otherwise, the number of
+    // modified/unmodified documents would be off by one depending on how our cluster is configured.
+    const docSize = 1024 * 1023;
     const numDocs = 20;
-    const maxDocsInBatch = maxBatchSize / docSize;
+    const maxDocsInBatch = Math.floor(maxBatchSize / docSize);
 
     assert(source.drop());
     dropWithoutImplicitRecreate(target.getName());
diff --git a/jstests/aggregation/sources/merge/mode_replace_fail.js b/jstests/aggregation/sources/merge/mode_replace_fail.js
index 88582238c8f..cc3df1b1ce7 100644
--- a/jstests/aggregation/sources/merge/mode_replace_fail.js
+++ b/jstests/aggregation/sources/merge/mode_replace_fail.js
@@ -90,9 +90,13 @@ const pipeline = [mergeStage];
 // and updated.
 (function testMergeUnorderedBatchUpdate() {
     const maxBatchSize = 16 * 1024 * 1024;  // 16MB
-    const docSize = 1024 * 1024;            // 1MB
+
+    // Each document is just under 1MB in order to allow for some extra space for writes that need
+    // to be serialized over the wire in certain cluster configurations. Otherwise, the number of
+    // modified/unmodified documents would be off by one depending on how our cluster is configured.
+    const docSize = 1024 * 1023;  // 1MB
     const numDocs = 20;
-    const maxDocsInBatch = maxBatchSize / docSize;
+    const maxDocsInBatch = Math.floor(maxBatchSize / docSize);
 
     assert(source.drop());
     dropWithoutImplicitRecreate(target.getName());