SERVER-28100 ReplSetTest.awaitNodesAgreeOnPrimary() accepts an optional expectedPrimaryIndex

(cherry-picked from commit 8d146d21661cf755fc15b7eefec147ce85ffde8c)

1 files changed, 6 insertions, 2 deletions

diff --git a/src/mongo/shell/replsettest.js b/src/mongo/shell/replsettest.js
index 59681b90f8c..667cef8c430 100644
--- a/src/mongo/shell/replsettest.js
+++ b/src/mongo/shell/replsettest.js
@@ -525,13 +525,17 @@ var ReplSetTest = function(opts) {
 
     /**
      * Blocks until all nodes agree on who the primary is.
+     * If 'expectedPrimaryNodeId' is provided, ensure that every node is seeing this node as the
+     * primary. Otherwise, ensure that all the nodes in the set agree with the first node on the
+     * identity of the primary.
      */
-    this.awaitNodesAgreeOnPrimary = function(timeout, nodes) {
+    this.awaitNodesAgreeOnPrimary = function(timeout, nodes, expectedPrimaryNodeId) {
         timeout = timeout || self.kDefaultTimeoutMS;
         nodes = nodes || self.nodes;
+        expectedPrimaryNodeId = expectedPrimaryNodeId || -1;
 
         assert.soonNoExcept(function() {
-            var primary = -1;
+            var primary = expectedPrimaryNodeId;
 
             for (var i = 0; i < nodes.length; i++) {
                 var replSetGetStatus = nodes[i].getDB("admin").runCommand({replSetGetStatus: 1});