SERVER-77202 create ReshardCollectionCmdTest util

3 files changed, 262 insertions, 202 deletions

diff --git a/jstests/sharding/libs/reshard_collection_util.js b/jstests/sharding/libs/reshard_collection_util.js
new file mode 100644
index 00000000000..d89e7ba7f51
--- /dev/null
+++ b/jstests/sharding/libs/reshard_collection_util.js
@@ -0,0 +1,224 @@
+/**
+ * Util class for testing reshardCollection cmd.
+ */
+
+"use strict";
+
+load("jstests/libs/uuid_util.js");
+
+class ReshardCollectionCmdTest {
+    constructor(testConfig) {
+        assert(testConfig.st && testConfig.dbName && testConfig.collName &&
+               testConfig.numInitialDocs);
+        this._st = testConfig.st;
+        this._mongos = this._st.s0;
+        this._mongosConfig = this._mongos.getDB('config');
+        this._dbName = testConfig.dbName;
+        this._collName = testConfig.collName;
+        this._ns = this._dbName + "." + this._collName;
+        this._numInitialDocs = testConfig.numInitialDocs;
+
+        this._shardToRSMap = {};
+        this._shardToRSMap[this._st.shard0.shardName] = this._st.rs0;
+        this._shardToRSMap[this._st.shard1.shardName] = this._st.rs1;
+        this._shardIdToShardMap = {};
+        this._shardIdToShardMap[this._st.shard0.shardName] = this._st.shard0;
+        this._shardIdToShardMap[this._st.shard1.shardName] = this._st.shard1;
+    }
+
+    _getUUIDFromCollectionInfo(dbName, collName, collInfo) {
+        if (collInfo) {
+            return extractUUIDFromObject(collInfo.info.uuid);
+        }
+
+        const uuidObject = getUUIDFromListCollections(this._mongos.getDB(dbName), collName);
+        return extractUUIDFromObject(uuidObject);
+    }
+
+    _constructTemporaryReshardingCollName(dbName, collName, collInfo) {
+        const existingUUID = this._getUUIDFromCollectionInfo(dbName, collName, collInfo);
+        return 'system.resharding.' + existingUUID;
+    }
+
+    _getAllShardIdsFromExpectedChunks(expectedChunks) {
+        let shardIds = new Set();
+        expectedChunks.forEach(chunk => {
+            shardIds.add(chunk.recipientShardId);
+        });
+        return shardIds;
+    }
+
+    _verifyChunksMatchExpected(numExpectedChunks, presetExpectedChunks) {
+        let collEntry =
+            this._mongos.getDB('config').getCollection('collections').findOne({_id: this._ns});
+        let chunkQuery = {uuid: collEntry.uuid};
+
+        const reshardedChunks = this._mongosConfig.chunks.find(chunkQuery).toArray();
+
+        if (presetExpectedChunks) {
+            presetExpectedChunks.sort();
+        }
+
+        reshardedChunks.sort();
+        assert.eq(numExpectedChunks, reshardedChunks.length, tojson(reshardedChunks));
+
+        let shardChunkCounts = {};
+        let incChunkCount = key => {
+            if (shardChunkCounts.hasOwnProperty(key)) {
+                shardChunkCounts[key]++;
+            } else {
+                shardChunkCounts[key] = 1;
+            }
+        };
+
+        for (let i = 0; i < numExpectedChunks; i++) {
+            incChunkCount(reshardedChunks[i].shard);
+
+            // match exact chunk boundaries for presetExpectedChunks
+            if (presetExpectedChunks) {
+                assert.eq(presetExpectedChunks[i].recipientShardId, reshardedChunks[i].shard);
+                assert.eq(presetExpectedChunks[i].min, reshardedChunks[i].min);
+                assert.eq(presetExpectedChunks[i].max, reshardedChunks[i].max);
+            }
+        }
+
+        // if presetChunks not specified, we only assert that chunks counts are balanced across
+        // shards
+        if (!presetExpectedChunks) {
+            let maxDiff = 0;
+            let shards = Object.keys(shardChunkCounts);
+
+            shards.forEach(shard1 => {
+                shards.forEach(shard2 => {
+                    let diff = Math.abs(shardChunkCounts[shard1] - shardChunkCounts[shard2]);
+                    maxDiff = (diff > maxDiff) ? diff : maxDiff;
+                });
+            });
+
+            assert.lte(maxDiff, 1, tojson(reshardedChunks));
+        }
+    }
+
+    _verifyCollectionExistenceForConn(collName, expectedToExist, conn) {
+        const doesExist = Boolean(conn.getDB(this._dbName)[collName].exists());
+        assert.eq(doesExist, expectedToExist);
+    }
+
+    _verifyTemporaryReshardingCollectionExistsWithCorrectOptions(expectedRecipientShards) {
+        const originalCollInfo =
+            this._mongos.getDB(this._dbName).getCollectionInfos({name: this._collName})[0];
+        assert.neq(originalCollInfo, undefined);
+
+        const tempReshardingCollName = this._constructTemporaryReshardingCollName(
+            this._dbName, this._collName, originalCollInfo);
+        this._verifyCollectionExistenceForConn(tempReshardingCollName, false, this._mongos);
+
+        expectedRecipientShards.forEach(shardId => {
+            const rsPrimary = this._shardToRSMap[shardId].getPrimary();
+            this._verifyCollectionExistenceForConn(this._collName, true, rsPrimary);
+            this._verifyCollectionExistenceForConn(tempReshardingCollName, false, rsPrimary);
+            ShardedIndexUtil.assertIndexExistsOnShard(
+                this._shardIdToShardMap[shardId], this._dbName, this._collName, {newKey: 1});
+        });
+    }
+
+    _verifyAllShardingCollectionsRemoved(tempReshardingCollName) {
+        assert.eq(0, this._mongos.getDB(this._dbName)[tempReshardingCollName].find().itcount());
+        assert.eq(0, this._mongosConfig.reshardingOperations.find({ns: this._ns}).itcount());
+        assert.eq(
+            0, this._mongosConfig.collections.find({reshardingFields: {$exists: true}}).itcount());
+        assert.eq(0,
+                  this._st.rs0.getPrimary()
+                      .getDB('config')
+                      .localReshardingOperations.donor.find({ns: this._ns})
+                      .itcount());
+        assert.eq(0,
+                  this._st.rs0.getPrimary()
+                      .getDB('config')
+                      .localReshardingOperations.recipient.find({ns: this._ns})
+                      .itcount());
+        assert.eq(0,
+                  this._st.rs1.getPrimary()
+                      .getDB('config')
+                      .localReshardingOperations.donor.find({ns: this._ns})
+                      .itcount());
+        assert.eq(0,
+                  this._st.rs1.getPrimary()
+                      .getDB('config')
+                      .localReshardingOperations.recipient.find({ns: this._ns})
+                      .itcount());
+    }
+
+    _verifyTagsDocumentsAfterOperationCompletes(ns, shardKeyPattern) {
+        const tagsArr = this._mongos.getCollection('config.tags').find({ns: ns}).toArray();
+        for (let i = 0; i < tagsArr.length; ++i) {
+            assert.eq(Object.keys(tagsArr[i]["min"]), shardKeyPattern);
+            assert.eq(Object.keys(tagsArr[i]["max"]), shardKeyPattern);
+        }
+    }
+
+    assertReshardCollOkWithPreset(commandObj, presetReshardedChunks) {
+        assert.commandWorked(
+            this._mongos.adminCommand({shardCollection: this._ns, key: {oldKey: 1}}));
+
+        let bulk = this._mongos.getDB(this._dbName)
+                       .getCollection(this._collName)
+                       .initializeOrderedBulkOp();
+        for (let x = 0; x < this._numInitialDocs; x++) {
+            bulk.insert({oldKey: x, newKey: this._numInitialDocs - x});
+        }
+        assert.commandWorked(bulk.execute());
+
+        commandObj._presetReshardedChunks = presetReshardedChunks;
+        const tempReshardingCollName =
+            this._constructTemporaryReshardingCollName(this._dbName, this._collName);
+
+        assert.commandWorked(this._mongos.adminCommand(commandObj));
+
+        this._verifyTemporaryReshardingCollectionExistsWithCorrectOptions(
+            this._getAllShardIdsFromExpectedChunks(presetReshardedChunks));
+
+        this._verifyTagsDocumentsAfterOperationCompletes(this._ns, Object.keys(commandObj.key));
+
+        this._verifyChunksMatchExpected(presetReshardedChunks.length, presetReshardedChunks);
+
+        this._mongos.getDB(this._dbName)[this._collName].drop();
+        this._verifyAllShardingCollectionsRemoved(tempReshardingCollName);
+    }
+
+    /**
+     * Run reshardCollection and check the number of chunks is as expected.
+     * @param {Object} commandObj The reshardCollection cmd to execute.
+     * @param {Number} expectedChunkNum Number of chunks to have after reshardCollection.
+     * @param {Object[]} expectedChunks Expected chunk distribution after reshardCollection.
+     */
+    assertReshardCollOk(commandObj, expectedChunkNum, expectedChunks) {
+        assert.commandWorked(
+            this._mongos.adminCommand({shardCollection: this._ns, key: {oldKey: 1}}));
+
+        let bulk = this._mongos.getDB(this._dbName)
+                       .getCollection(this._collName)
+                       .initializeOrderedBulkOp();
+        for (let x = 0; x < this._numInitialDocs; x++) {
+            bulk.insert({oldKey: x, newKey: this._numInitialDocs - x});
+        }
+        assert.commandWorked(bulk.execute());
+
+        const tempReshardingCollName =
+            this._constructTemporaryReshardingCollName(this._dbName, this._collName);
+
+        assert.commandWorked(this._mongos.adminCommand(commandObj));
+
+        if (expectedChunks) {
+            this._verifyTemporaryReshardingCollectionExistsWithCorrectOptions(
+                this._getAllShardIdsFromExpectedChunks(expectedChunks));
+        }
+
+        this._verifyTagsDocumentsAfterOperationCompletes(this._ns, Object.keys(commandObj.key));
+
+        this._verifyChunksMatchExpected(expectedChunkNum, expectedChunks);
+
+        this._mongos.getDB(this._dbName)[this._collName].drop();
+        this._verifyAllShardingCollectionsRemoved(tempReshardingCollName);
+    }
+}
\ No newline at end of file
diff --git a/jstests/sharding/reshard_collection_basic.js b/jstests/sharding/reshard_collection_basic.js
index 9ad99cfa753..d1f2b874d45 100644
--- a/jstests/sharding/reshard_collection_basic.js
+++ b/jstests/sharding/reshard_collection_basic.js
@@ -6,9 +6,9 @@
 //
 
 load("jstests/libs/fail_point_util.js");
-load("jstests/libs/uuid_util.js");
 load("jstests/sharding/libs/find_chunks_util.js");
 load("jstests/libs/discover_topology.js");
+load("jstests/sharding/libs/reshard_collection_util.js");
 
 (function() {
 'use strict';
@@ -18,8 +18,9 @@ const kDbName = 'db';
 const collName = 'foo';
 const ns = kDbName + '.' + collName;
 const mongos = st.s0;
-const mongosConfig = mongos.getDB('config');
 const kNumInitialDocs = 500;
+const reshardCmdTest =
+    new ReshardCollectionCmdTest({st, dbName: kDbName, collName, numInitialDocs: kNumInitialDocs});
 
 const criticalSectionTimeoutMS = 24 * 60 * 60 * 1000; /* 1 day */
 const topology = DiscoverTopology.findConnectedNodes(mongos);
@@ -27,183 +28,6 @@ const coordinator = new Mongo(topology.configsvr.nodes[0]);
 assert.commandWorked(coordinator.getDB("admin").adminCommand(
     {setParameter: 1, reshardingCriticalSectionTimeoutMillis: criticalSectionTimeoutMS}));
 
-let shardToRSMap = {};
-shardToRSMap[st.shard0.shardName] = st.rs0;
-shardToRSMap[st.shard1.shardName] = st.rs1;
-
-let shardIdToShardMap = {};
-shardIdToShardMap[st.shard0.shardName] = st.shard0;
-shardIdToShardMap[st.shard1.shardName] = st.shard1;
-
-let getUUIDFromCollectionInfo = (dbName, collName, collInfo) => {
-    if (collInfo) {
-        return extractUUIDFromObject(collInfo.info.uuid);
-    }
-
-    const uuidObject = getUUIDFromListCollections(mongos.getDB(dbName), collName);
-    return extractUUIDFromObject(uuidObject);
-};
-
-let constructTemporaryReshardingCollName = (dbName, collName, collInfo) => {
-    const existingUUID = getUUIDFromCollectionInfo(dbName, collName, collInfo);
-    return 'system.resharding.' + existingUUID;
-};
-
-let getAllShardIdsFromExpectedChunks = (expectedChunks) => {
-    let shardIds = new Set();
-    expectedChunks.forEach(chunk => {
-        shardIds.add(chunk.recipientShardId);
-    });
-    return shardIds;
-};
-
-let verifyChunksMatchExpected = (numExpectedChunks, presetExpectedChunks) => {
-    let collEntry = mongos.getDB('config').getCollection('collections').findOne({_id: ns});
-    let chunkQuery = {uuid: collEntry.uuid};
-
-    const reshardedChunks = mongosConfig.chunks.find(chunkQuery).toArray();
-
-    if (presetExpectedChunks) {
-        presetExpectedChunks.sort();
-    }
-
-    reshardedChunks.sort();
-    assert.eq(numExpectedChunks, reshardedChunks.length, tojson(reshardedChunks));
-
-    let shardChunkCounts = {};
-    let incChunkCount = key => {
-        if (shardChunkCounts.hasOwnProperty(key)) {
-            shardChunkCounts[key]++;
-        } else {
-            shardChunkCounts[key] = 1;
-        }
-    };
-
-    for (let i = 0; i < numExpectedChunks; i++) {
-        incChunkCount(reshardedChunks[i].shard);
-
-        // match exact chunk boundaries for presetExpectedChunks
-        if (presetExpectedChunks) {
-            assert.eq(presetExpectedChunks[i].recipientShardId, reshardedChunks[i].shard);
-            assert.eq(presetExpectedChunks[i].min, reshardedChunks[i].min);
-            assert.eq(presetExpectedChunks[i].max, reshardedChunks[i].max);
-        }
-    }
-
-    // if presetChunks not specified, we only assert that chunks counts are balanced across shards
-    if (!presetExpectedChunks) {
-        let maxDiff = 0;
-        let shards = Object.keys(shardChunkCounts);
-
-        shards.forEach(shard1 => {
-            shards.forEach(shard2 => {
-                let diff = Math.abs(shardChunkCounts[shard1] - shardChunkCounts[shard2]);
-                maxDiff = (diff > maxDiff) ? diff : maxDiff;
-            });
-        });
-
-        assert.lte(maxDiff, 1, tojson(reshardedChunks));
-    }
-};
-
-let verifyCollectionExistenceForConn = (collName, expectedToExist, conn) => {
-    const doesExist = Boolean(conn.getDB(kDbName)[collName].exists());
-    assert.eq(doesExist, expectedToExist);
-};
-
-let verifyTemporaryReshardingCollectionExistsWithCorrectOptions = (expectedRecipientShards) => {
-    const originalCollInfo = mongos.getDB(kDbName).getCollectionInfos({name: collName})[0];
-    assert.neq(originalCollInfo, undefined);
-
-    const tempReshardingCollName =
-        constructTemporaryReshardingCollName(kDbName, collName, originalCollInfo);
-    verifyCollectionExistenceForConn(tempReshardingCollName, false, mongos);
-
-    expectedRecipientShards.forEach(shardId => {
-        const rsPrimary = shardToRSMap[shardId].getPrimary();
-        verifyCollectionExistenceForConn(collName, true, rsPrimary);
-        verifyCollectionExistenceForConn(tempReshardingCollName, false, rsPrimary);
-        ShardedIndexUtil.assertIndexExistsOnShard(
-            shardIdToShardMap[shardId], kDbName, collName, {newKey: 1});
-    });
-};
-
-let verifyAllShardingCollectionsRemoved = (tempReshardingCollName) => {
-    assert.eq(0, mongos.getDB(kDbName)[tempReshardingCollName].find().itcount());
-    assert.eq(0, mongosConfig.reshardingOperations.find({ns}).itcount());
-    assert.eq(0, mongosConfig.collections.find({reshardingFields: {$exists: true}}).itcount());
-    assert.eq(
-        0,
-        st.rs0.getPrimary().getDB('config').localReshardingOperations.donor.find({ns}).itcount());
-    assert.eq(0,
-              st.rs0.getPrimary()
-                  .getDB('config')
-                  .localReshardingOperations.recipient.find({ns})
-                  .itcount());
-    assert.eq(
-        0,
-        st.rs1.getPrimary().getDB('config').localReshardingOperations.donor.find({ns}).itcount());
-    assert.eq(0,
-              st.rs1.getPrimary()
-                  .getDB('config')
-                  .localReshardingOperations.recipient.find({ns})
-                  .itcount());
-};
-
-let verifyTagsDocumentsAfterOperationCompletes = (ns, shardKeyPattern) => {
-    const tagsArr = mongos.getCollection('config.tags').find({ns: ns}).toArray();
-    for (let i = 0; i < tagsArr.length; ++i) {
-        assert.eq(Object.keys(tagsArr[i]["min"]), shardKeyPattern);
-        assert.eq(Object.keys(tagsArr[i]["max"]), shardKeyPattern);
-    }
-};
-
-let assertReshardCollOkWithPreset = (commandObj, presetReshardedChunks) => {
-    assert.commandWorked(mongos.adminCommand({shardCollection: ns, key: {oldKey: 1}}));
-
-    let bulk = mongos.getDB(kDbName).getCollection(collName).initializeOrderedBulkOp();
-    for (let x = 0; x < kNumInitialDocs; x++) {
-        bulk.insert({oldKey: x, newKey: kNumInitialDocs - x});
-    }
-    assert.commandWorked(bulk.execute());
-
-    commandObj._presetReshardedChunks = presetReshardedChunks;
-    const tempReshardingCollName = constructTemporaryReshardingCollName(kDbName, collName);
-
-    assert.commandWorked(mongos.adminCommand(commandObj));
-
-    verifyTemporaryReshardingCollectionExistsWithCorrectOptions(
-        getAllShardIdsFromExpectedChunks(presetReshardedChunks));
-
-    verifyTagsDocumentsAfterOperationCompletes(ns, Object.keys(commandObj.key));
-
-    verifyChunksMatchExpected(presetReshardedChunks.length, presetReshardedChunks);
-
-    mongos.getDB(kDbName)[collName].drop();
-    verifyAllShardingCollectionsRemoved(tempReshardingCollName);
-};
-
-let assertReshardCollOk = (commandObj, expectedChunks) => {
-    assert.commandWorked(mongos.adminCommand({shardCollection: ns, key: {oldKey: 1}}));
-
-    let bulk = mongos.getDB(kDbName).getCollection(collName).initializeOrderedBulkOp();
-    for (let x = 0; x < kNumInitialDocs; x++) {
-        bulk.insert({oldKey: x, newKey: kNumInitialDocs - x});
-    }
-    assert.commandWorked(bulk.execute());
-
-    const tempReshardingCollName = constructTemporaryReshardingCollName(kDbName, collName);
-
-    assert.commandWorked(mongos.adminCommand(commandObj));
-
-    verifyTagsDocumentsAfterOperationCompletes(ns, Object.keys(commandObj.key));
-
-    verifyChunksMatchExpected(expectedChunks);
-
-    mongos.getDB(kDbName)[collName].drop();
-    verifyAllShardingCollectionsRemoved(tempReshardingCollName);
-};
-
 let presetReshardedChunks =
     [{recipientShardId: st.shard1.shardName, min: {newKey: MinKey}, max: {newKey: MaxKey}}];
 
@@ -293,17 +117,19 @@ assert.commandFailedWithCode(mongos.getDB('test').system.resharding.mycoll.inser
 mongos.getDB(kDbName)[collName].drop();
 
 jsTest.log("Succeed when correct locale is provided.");
-assertReshardCollOk({reshardCollection: ns, key: {newKey: 1}, collation: {locale: 'simple'}}, 1);
+reshardCmdTest.assertReshardCollOk(
+    {reshardCollection: ns, key: {newKey: 1}, collation: {locale: 'simple'}}, 1);
 
 jsTest.log("Succeed base case.");
-assertReshardCollOk({reshardCollection: ns, key: {newKey: 1}}, 1);
+reshardCmdTest.assertReshardCollOk({reshardCollection: ns, key: {newKey: 1}}, 1);
 
 jsTest.log("Succeed if unique is specified and is false.");
-assertReshardCollOk({reshardCollection: ns, key: {newKey: 1}, unique: false}, 1);
+reshardCmdTest.assertReshardCollOk({reshardCollection: ns, key: {newKey: 1}, unique: false}, 1);
 
 jsTest.log(
     "Succeed if _presetReshardedChunks is provided and test commands are enabled (default).");
-assertReshardCollOkWithPreset({reshardCollection: ns, key: {newKey: 1}}, presetReshardedChunks);
+reshardCmdTest.assertReshardCollOkWithPreset({reshardCollection: ns, key: {newKey: 1}},
+                                             presetReshardedChunks);
 
 presetReshardedChunks = [
     {recipientShardId: st.shard0.shardName, min: {newKey: MinKey}, max: {newKey: 0}},
@@ -311,19 +137,19 @@ presetReshardedChunks = [
 ];
 
 jsTest.log("Succeed if all optional fields and numInitialChunks are provided with correct values.");
-assertReshardCollOk({
+reshardCmdTest.assertReshardCollOk({
     reshardCollection: ns,
     key: {newKey: 1},
     unique: false,
     collation: {locale: 'simple'},
     numInitialChunks: 2,
 },
-                    2);
+                                   2);
 
 jsTest.log(
     "Succeed if all optional fields and _presetReshardedChunks are provided with correct values" +
     " and test commands are enabled (default).");
-assertReshardCollOkWithPreset(
+reshardCmdTest.assertReshardCollOkWithPreset(
     {reshardCollection: ns, key: {newKey: 1}, unique: false, collation: {locale: 'simple'}},
     presetReshardedChunks);
 
@@ -331,17 +157,17 @@ jsTest.log("Succeed if the zone provided is assigned to a shard but not a range
            " collection.");
 const newZoneName = 'x2';
 assert.commandWorked(st.s.adminCommand({addShardToZone: st.shard1.shardName, zone: newZoneName}));
-assertReshardCollOk({
+reshardCmdTest.assertReshardCollOk({
     reshardCollection: ns,
     key: {newKey: 1},
     unique: false,
     collation: {locale: 'simple'},
     zones: [{zone: newZoneName, min: {newKey: 5}, max: {newKey: 10}}]
 },
-                    3);
+                                   3);
 
 jsTest.log("Succeed if resulting chunks all end up in one shard.");
-assertReshardCollOk({
+reshardCmdTest.assertReshardCollOk({
     reshardCollection: ns,
     key: {newKey: 1},
     unique: false,
@@ -349,10 +175,10 @@ assertReshardCollOk({
     collation: {locale: 'simple'},
     zones: [{zone: newZoneName, min: {newKey: MinKey}, max: {newKey: MaxKey}}]
 },
-                    1);
+                                   1);
 
 jsTest.log("Succeed if zones are empty");
-assertReshardCollOk({
+reshardCmdTest.assertReshardCollOk({
     reshardCollection: ns,
     key: {newKey: 1},
     unique: false,
@@ -360,21 +186,21 @@ assertReshardCollOk({
     collation: {locale: 'simple'},
     zones: []
 },
-                    1);
+                                   1);
 
 jsTest.log("Succeed if zones are not empty.");
 assert.commandWorked(
     mongos.adminCommand({addShardToZone: st.shard1.shardName, zone: existingZoneName}));
 assert.commandWorked(st.s.adminCommand(
     {updateZoneKeyRange: ns, min: {oldKey: 0}, max: {oldKey: 5}, zone: existingZoneName}));
-assertReshardCollOk({
+reshardCmdTest.assertReshardCollOk({
     reshardCollection: ns,
     key: {oldKey: 1, newKey: 1},
     unique: false,
     collation: {locale: 'simple'},
     zones: [{zone: existingZoneName, min: {oldKey: 0}, max: {oldKey: 5}}]
 },
-                    3);
+                                   3);
 
 jsTest.log("Succeed with hashed shard key that provides enough cardinality.");
 assert.commandWorked(
diff --git a/jstests/sharding/reshard_collection_resharding_improvements_basic.js b/jstests/sharding/reshard_collection_resharding_improvements_basic.js
index 3244a3cac0d..9339ef9bddb 100644
--- a/jstests/sharding/reshard_collection_resharding_improvements_basic.js
+++ b/jstests/sharding/reshard_collection_resharding_improvements_basic.js
@@ -11,12 +11,16 @@
 'use strict';
 
 load("jstests/libs/feature_flag_util.js");
+load("jstests/sharding/libs/reshard_collection_util.js");
 
 const st = new ShardingTest({mongos: 1, shards: 2});
 const kDbName = 'db';
 const collName = 'foo';
 const ns = kDbName + '.' + collName;
 const mongos = st.s0;
+const kNumInitialDocs = 500;
+const reshardCmdTest =
+    new ReshardCollectionCmdTest({st, dbName: kDbName, collName, numInitialDocs: kNumInitialDocs});
 
 const testShardDistribution = (mongos) => {
     if (!FeatureFlagUtil.isEnabled(mongos, "ReshardingImprovements")) {
@@ -105,14 +109,20 @@ const testShardDistribution = (mongos) => {
         shardDistribution: [{shard: st.shard0.shardName}, {shard: st.shard1.shardName}]
     }),
                                  ErrorCodes.InvalidOptions);
-    assert.commandWorked(mongos.adminCommand({
-        reshardCollection: ns,
-        key: {newKey: 1},
-        shardDistribution: [
-            {shard: st.shard0.shardName, min: {newKey: MinKey}, max: {newKey: 0}},
-            {shard: st.shard1.shardName, min: {newKey: 0}, max: {newKey: MaxKey}}
-        ]
-    }));
+    reshardCmdTest.assertReshardCollOk(
+        {
+            reshardCollection: ns,
+            key: {newKey: 1},
+            shardDistribution: [
+                {shard: st.shard0.shardName, min: {newKey: MinKey}, max: {newKey: 0}},
+                {shard: st.shard1.shardName, min: {newKey: 0}, max: {newKey: MaxKey}}
+            ]
+        },
+        2,
+        [
+            {recipientShardId: st.shard0.shardName, min: {newKey: MinKey}, max: {newKey: 0}},
+            {recipientShardId: st.shard1.shardName, min: {newKey: 0}, max: {newKey: MaxKey}}
+        ]);
 };
 
 testShardDistribution(mongos);