SERVER-44602 Test that chunks and documents are moved after zone changes

1 files changed, 88 insertions, 0 deletions

diff --git a/jstests/sharding/libs/zone_changes_util.js b/jstests/sharding/libs/zone_changes_util.js
new file mode 100644
index 00000000000..f45a5e28ef7
--- /dev/null
+++ b/jstests/sharding/libs/zone_changes_util.js
@@ -0,0 +1,88 @@
+load("jstests/sharding/libs/chunk_bounds_util.js");
+
+/**
+ * Asserts that the given shards have the given chunks.
+ *
+ * @param shardChunkBounds {Object} a map from each shard name to an array of the bounds for all
+ *                                  the chunks on the shard. Each pair of chunk bounds is an array
+ *                                  of the form [minKey, maxKey].
+ */
+function assertChunksOnShards(configDB, ns, shardChunkBounds) {
+    for (let [shardName, chunkBounds] of Object.entries(shardChunkBounds)) {
+        for (let bounds of chunkBounds) {
+            assert.eq(
+                shardName,
+                configDB.chunks.findOne({ns: ns, min: bounds[0], max: bounds[1]}).shard,
+                "expected to find chunk " + tojson(bounds) + " on shard \"" + shardName + "\"");
+        }
+    }
+}
+
+/**
+ * Asserts that the docs are on the shards that own their corresponding chunks.
+ *
+ * @param shardChunkBounds {Object} a map from each shard name to an array of the bounds for all
+ *                                  the chunks on the shard. Each pair of chunk bounds is an array
+ *                                  of the form [minKey, maxKey].
+ * @param shardKey         {Object} a map from each shard key field to 1 if the collection uses
+ *                                  range based sharding and "hashed" if the collection uses
+ *                                  hashed sharding. (i.e. equivalent to the value passed for the
+ *                                  "key" field for the shardCollection command).
+ */
+function assertDocsOnShards(st, ns, shardChunkBounds, docs, shardKey) {
+    for (let doc of docs) {
+        let docShardKey = {};
+        for (const [k, v] of Object.entries(shardKey)) {
+            docShardKey[k] = (v == "hashed") ? convertShardKeyToHashed(doc[k]) : doc[k];
+        }
+        let shard = chunkBoundsUtil.findShardForShardKey(st, shardChunkBounds, docShardKey);
+        assert.eq(1,
+                  shard.getCollection(ns).count(doc),
+                  "expected to find doc " + tojson(doc) + " on shard \"" + shard.shardName + "\"");
+    }
+}
+
+/**
+ * Asserts that the given shards have the given tags.
+ *
+ * @param shardTags {Object} a map from each shard name to an array of strings representing the zone
+ *                           names that the shard owns.
+ */
+function assertShardTags(configDB, shardTags) {
+    for (let [shardName, tags] of Object.entries(shardTags)) {
+        assert.eq(tags.sort(),
+                  configDB.shards.findOne({_id: shardName}).tags.sort(),
+                  "expected shard \"" + shardName + "\" to have tags " + tojson(tags.sort()));
+    }
+}
+
+/**
+ * Adds toShard to zone and removes fromShard from zone.
+ */
+function moveZoneToShard(st, zoneName, fromShard, toShard) {
+    assert.commandWorked(st.s.adminCommand({addShardToZone: toShard.shardName, zone: zoneName}));
+    assert.commandWorked(
+        st.s.adminCommand({removeShardFromZone: fromShard.shardName, zone: zoneName}));
+}
+
+/**
+ * Starts the balancer, lets it run for the given number of rounds, then stops the
+ * balancer.
+ */
+function runBalancer(st, numRounds) {
+    st.startBalancer();
+    for (let i = 0; i < numRounds; i++) {
+        st.awaitBalancerRound();
+    }
+    st.stopBalancer();
+}
+
+/**
+ * Updates the zone key range for the given namespace.
+ */
+function updateZoneKeyRange(st, ns, zoneName, fromRange, toRange) {
+    assert.commandWorked(st.s.adminCommand(
+        {updateZoneKeyRange: ns, min: fromRange[0], max: fromRange[1], zone: null}));
+    assert.commandWorked(st.s.adminCommand(
+        {updateZoneKeyRange: ns, min: toRange[0], max: toRange[1], zone: zoneName}));
+}