1 files changed, 46 insertions, 19 deletions
diff --git a/jstests/sharding/shard_collection_verify_initial_chunks.js b/jstests/sharding/shard_collection_verify_initial_chunks.js
index 0538bed1b34..e7072132b11 100644
--- a/jstests/sharding/shard_collection_verify_initial_chunks.js
+++ b/jstests/sharding/shard_collection_verify_initial_chunks.js
@@ -1,30 +1,57 @@
-//
-// Verify numInitialChunks can not be set for non hashed key or nonempty collections
-//
-
+/**
+ * Verify initial chunks are properly created and distributed in various combinations of shard key
+ * and empty/non-empty collections.
+ */
 (function() {
     'use strict';
 
-    var st = new ShardingTest({mongos: 1, shards: 2});
-    var kDbName = 'db';
-    var mongos = st.s0;
+    let st = new ShardingTest({mongos: 1, shards: 2});
+    let mongos = st.s0;
 
-    assert.commandWorked(mongos.adminCommand({enableSharding: kDbName}));
+    let config = mongos.getDB("config");
+    let db = mongos.getDB('TestDB');
 
-    assert.commandFailed(mongos.adminCommand(
-        {shardCollection: kDbName + '.foo', key: {aKey: 1}, numInitialChunks: 5}));
+    assert.commandWorked(mongos.adminCommand({enableSharding: 'TestDB'}));
+    st.ensurePrimaryShard('TestDB', st.shard1.shardName);
 
-    assert.writeOK(mongos.getDB(kDbName).foo.insert({aKey: 1}));
-    assert.commandWorked(mongos.getDB(kDbName).foo.createIndex({aKey: "hashed"}));
-    assert.commandFailed(mongos.adminCommand(
-        {shardCollection: kDbName + '.foo', key: {aKey: "hashed"}, numInitialChunks: 5}));
+    function checkChunkCounts(collName, chunksOnShard0, chunksOnShard1) {
+        let counts = st.chunkCounts(collName, 'TestDB');
+        assert.eq(chunksOnShard0,
+                  counts[st.shard0.shardName],
+                  'Count mismatch on shard0: ' + tojson(counts));
+        assert.eq(chunksOnShard1,
+                  counts[st.shard1.shardName],
+                  'Count mismatch on shard1: ' + tojson(counts));
+    }
 
-    assert.writeOK(mongos.getDB(kDbName).foo.remove({}));
+    // Unsupported: Range sharding + numInitialChunks
+    assert.commandFailed(mongos.adminCommand(
+        {shardCollection: 'TestDB.RangeCollEmpty', key: {aKey: 1}, numInitialChunks: 6}));
+
+    // Unsupported: Hashed sharding + numInitialChunks + non-empty collection
+    assert.writeOK(db.HashedCollNotEmpty.insert({aKey: 1}));
+    assert.commandWorked(db.HashedCollNotEmpty.createIndex({aKey: "hashed"}));
+    assert.commandFailed(mongos.adminCommand({
+        shardCollection: 'TestDB.HashedCollNotEmpty',
+        key: {aKey: "hashed"},
+        numInitialChunks: 6
+    }));
+
+    // Supported: Hashed sharding + numInitialChunks + empty collection
+    // Expected: Even chunk distribution
+    assert.commandWorked(db.HashedCollEmpty.createIndex({aKey: "hashed"}));
     assert.commandWorked(mongos.adminCommand(
-        {shardCollection: kDbName + '.foo', key: {aKey: "hashed"}, numInitialChunks: 5}));
-
-    mongos.getDB(kDbName).dropDatabase();
+        {shardCollection: 'TestDB.HashedCollEmpty', key: {aKey: "hashed"}, numInitialChunks: 6}));
+    checkChunkCounts('HashedCollEmpty', 3, 3);
+
+    // Supported: Hashed sharding + numInitialChunks + non-existent collection
+    // Expected: Even chunk distribution
+    assert.commandWorked(mongos.adminCommand({
+        shardCollection: 'TestDB.HashedCollNonExistent',
+        key: {aKey: "hashed"},
+        numInitialChunks: 6
+    }));
+    checkChunkCounts('HashedCollNonExistent', 3, 3);
 
     st.stop();
-
 })();