6 files changed, 26 insertions, 7 deletions
diff --git a/jstests/noPassthroughWithMongod/no_balance_collection.js b/jstests/noPassthroughWithMongod/no_balance_collection.js
index 3a856fe35b4..b548ce13cd9 100644
--- a/jstests/noPassthroughWithMongod/no_balance_collection.js
+++ b/jstests/noPassthroughWithMongod/no_balance_collection.js
@@ -72,7 +72,7 @@ sh.disableBalancing(collB);
 // Wait for the balancer to fully finish the last migration and write the changelog
 // MUST set db var here, ugly but necessary
 db = st.s0.getDB("config");
-st.awaitBalancerRound();
+st.waitForBalancer(true, 60000);
 
 // Make sure auto-migrates on insert don't move chunks
 var lastMigration = sh._lastMigration(collB);
diff --git a/jstests/sharding/auto_rebalance_parallel.js b/jstests/sharding/auto_rebalance_parallel.js
index c27fef1fa31..bb86c1fb9f1 100644
--- a/jstests/sharding/auto_rebalance_parallel.js
+++ b/jstests/sharding/auto_rebalance_parallel.js
@@ -54,8 +54,8 @@
     const testColl2InitialMoves = countMoves('TestDB.TestColl2');
 
     st.startBalancer();
-    st.awaitBalancerRound();
-    st.awaitBalancerRound();
+    st.waitForBalancer(true, 60000);
+    st.waitForBalancer(true, 60000);
     st.stopBalancer();
 
     checkCollectionBalanced('TestDB.TestColl1');
diff --git a/jstests/sharding/auto_rebalance_parallel_replica_sets.js b/jstests/sharding/auto_rebalance_parallel_replica_sets.js
index 965219c56ff..fdcd6457fe4 100644
--- a/jstests/sharding/auto_rebalance_parallel_replica_sets.js
+++ b/jstests/sharding/auto_rebalance_parallel_replica_sets.js
@@ -41,7 +41,7 @@
 
     // Do enable the balancer and wait for a single balancer round
     st.startBalancer();
-    st.awaitBalancerRound();
+    st.waitForBalancer(true, 60000);
     st.stopBalancer();
 
     assert.eq(1,
diff --git a/jstests/sharding/balancer_window.js b/jstests/sharding/balancer_window.js
index a60654b7d89..422085a537a 100644
--- a/jstests/sharding/balancer_window.js
+++ b/jstests/sharding/balancer_window.js
@@ -68,7 +68,7 @@
                                             true));
     st.startBalancer();
 
-    st.awaitBalancerRound();
+    st.waitForBalancer(true, 60000);
 
     var shard0ChunksAfter =
         configDB.chunks.find({ns: 'test.user', shard: st.shard0.shardName}).count();
@@ -84,7 +84,7 @@
         },
         true));
 
-    st.awaitBalancerRound();
+    st.waitForBalancer(true, 60000);
 
     shard0ChunksAfter = configDB.chunks.find({ns: 'test.user', shard: st.shard0.shardName}).count();
     assert.neq(shard0Chunks, shard0ChunksAfter);
diff --git a/jstests/sharding/enforce_zone_policy.js b/jstests/sharding/enforce_zone_policy.js
index 7aac1aa02bf..4c57885ce10 100644
--- a/jstests/sharding/enforce_zone_policy.js
+++ b/jstests/sharding/enforce_zone_policy.js
@@ -34,7 +34,7 @@
         assert.soon(
             checkFunc, 'Balance at step ' + stepName + ' did not happen', 3 * 60 * 1000, 2000);
 
-        st.awaitBalancerRound();
+        st.waitForBalancer(true, 60000);
         st.printShardingStatus(true);
         assert(checkFunc());
 
diff --git a/src/mongo/shell/utils_sh.js b/src/mongo/shell/utils_sh.js
index b872da66a2c..c83e2cff3ff 100644
--- a/src/mongo/shell/utils_sh.js
+++ b/src/mongo/shell/utils_sh.js
@@ -24,6 +24,10 @@ sh._getConfigDB = function() {
     return db.getSiblingDB("config");
 };
 
+sh._getBalancerStatus = function() {
+    return assert.commandWorked(sh._getConfigDB().adminCommand({balancerStatus: 1}));
+};
+
 sh._dataFormat = function(bytes) {
     if (bytes < 1024)
         return Math.floor(bytes) + "B";
@@ -233,6 +237,21 @@ sh.waitForPingChange = function(activePings, timeout, interval) {
     return remainingPings;
 };
 
+sh.waitForBalancer = function(wait, timeout, interval) {
+    if (typeof(wait) === 'undefined') {
+        wait = false;
+    }
+    var initialStatus = sh._getBalancerStatus();
+    if (!initialStatus.inBalancerRound && !wait) {
+        return;
+    }
+    var currentStatus;
+    assert.soon(function() {
+        currentStatus = sh._getBalancerStatus();
+        return (currentStatus.numBalancerRounds - initialStatus.numBalancerRounds) != 0;
+    }, 'Latest balancer status: ' + tojson(currentStatus), timeout, interval);
+};
+
 sh.disableBalancing = function(coll) {
     if (coll === undefined) {
         throw Error("Must specify collection");