1 files changed, 35 insertions, 15 deletions

diff --git a/jstests/sharding/change_streams_unsharded_becomes_sharded.js b/jstests/sharding/change_streams_unsharded_becomes_sharded.js
index e172b28f463..c59ae695486 100644
--- a/jstests/sharding/change_streams_unsharded_becomes_sharded.js
+++ b/jstests/sharding/change_streams_unsharded_becomes_sharded.js
@@ -1,5 +1,6 @@
 // Tests the behavior of change streams on a collection that was initially unsharded but then
-// becomes sharded.
+// becomes sharded. In particular, test that post-shardCollection inserts update their cached
+// 'documentKey' to include the new shard key.
 (function() {
     "use strict";
 
@@ -29,6 +30,7 @@
     const mongosColl = mongosDB[testName];
 
     mongosDB.createCollection(testName);
+    mongosColl.createIndex({x: 1});
 
     st.ensurePrimaryShard(mongosDB.getName(), st.rs0.getURL());
 
@@ -38,44 +40,62 @@
         cst.startWatchingChanges({pipeline: [{$changeStream: {}}], collection: mongosColl});
     assert.eq(0, cursor.firstBatch.length, "Cursor had changes: " + tojson(cursor));
 
+    // Verify that the cursor picks up documents inserted while the collection is unsharded. The
+    // 'documentKey' at this point is simply the _id field.
+    assert.writeOK(mongosColl.insert({_id: 0, x: 0}));
+    cst.assertNextChangesEqual({
+        cursor: cursor,
+        expectedChanges: [{
+            documentKey: {_id: 0},
+            fullDocument: {_id: 0, x: 0},
+            ns: {db: mongosDB.getName(), coll: mongosColl.getName()},
+            operationType: "insert",
+        }]
+    });
+
     // Enable sharding on the previously unsharded collection.
     assert.commandWorked(mongosDB.adminCommand({enableSharding: mongosDB.getName()}));
 
-    // Shard the collection on _id.
+    // Shard the collection on x.
     assert.commandWorked(
-        mongosDB.adminCommand({shardCollection: mongosColl.getFullName(), key: {_id: 1}}));
+        mongosDB.adminCommand({shardCollection: mongosColl.getFullName(), key: {x: 1}}));
+
+    // Ensure that the primary shard has an up-to-date routing table.
+    assert.commandWorked(st.rs0.getPrimary().getDB("admin").runCommand(
+        {_flushRoutingTableCacheUpdates: mongosColl.getFullName()}));
 
     // Split the collection into 2 chunks: [MinKey, 0), [0, MaxKey).
-    assert.commandWorked(
-        mongosDB.adminCommand({split: mongosColl.getFullName(), middle: {_id: 0}}));
+    assert.commandWorked(mongosDB.adminCommand({split: mongosColl.getFullName(), middle: {x: 0}}));
 
-    // Verify that the cursor is still valid and picks up the inserted document.
-    assert.writeOK(mongosColl.insert({_id: 1}));
+    // Verify that the cursor on the original shard is still valid and sees new inserted documents.
+    // The 'documentKey' field should now include the shard key, even before a 'kNewShardDetected'
+    // operation has been generated by the migration of a chunk to a new shard.
+    assert.writeOK(mongosColl.insert({_id: 1, x: 1}));
     cst.assertNextChangesEqual({
         cursor: cursor,
         expectedChanges: [{
-            documentKey: {_id: 1},
-            fullDocument: {_id: 1},
+            documentKey: {x: 1, _id: 1},
+            fullDocument: {_id: 1, x: 1},
             ns: {db: mongosDB.getName(), coll: mongosColl.getName()},
             operationType: "insert",
         }]
     });
 
-    // Move the [minKey, 0) chunk to shard1 and write a document to it.
+    // Move the [minKey, 0) chunk to shard1.
     assert.commandWorked(mongosDB.adminCommand({
         moveChunk: mongosColl.getFullName(),
-        find: {_id: -1},
+        find: {x: -1},
         to: st.rs1.getURL(),
         _waitForDelete: true
     }));
-    assert.writeOK(mongosColl.insert({_id: -1}));
 
-    // Make sure the change stream cursor sees the inserted document even after the moveChunk.
+    // Make sure the change stream cursor sees a document inserted on the recipient shard.
+    assert.writeOK(mongosColl.insert({_id: -1, x: -1}));
     cst.assertNextChangesEqual({
         cursor: cursor,
         expectedChanges: [{
-            documentKey: {_id: -1},
-            fullDocument: {_id: -1},
+            documentKey: {x: -1, _id: -1},
+            fullDocument: {_id: -1, x: -1},
             ns: {db: mongosDB.getName(), coll: mongosColl.getName()},
             operationType: "insert",
         }]