From 7ff4898e9703ae2eca368ab768f45bedc735e1a6 Mon Sep 17 00:00:00 2001
From: Bernie Hackett <bernie@10gen.com>
Date: Mon, 6 Feb 2017 13:30:49 -0800
Subject: [PATCH] Improve test suite primary discovery

---
 test/__init__.py | 32 ++++++++++++++++++++------------
 1 file changed, 20 insertions(+), 12 deletions(-)

diff --git a/test/__init__.py b/test/__init__.py
index d8f7aa05c..b66b3d9b4 100644
--- a/test/__init__.py
+++ b/test/__init__.py
@@ -17,6 +17,7 @@
 
 import os
 import sys
+import time
 import warnings
 
 import pymongo
@@ -49,18 +50,25 @@ def _split_host_port(host_port):
 # When testing against a replica set, 'host' and 'port' need
 # to represent the primary.
 try:
-    _ismaster = pymongo.MongoClient(host, port).admin.command('ismaster')
-    # Master slave topologies don't use setName. The members also aren't
-    # aware of one another (e.g. ismaster doesn't have a 'hosts' field).
-    if 'setName' in _ismaster:
-        _primary = _ismaster['primary']
-        host, port = _split_host_port(_primary)
-        _secondaries = set(_ismaster['hosts']) - set([_primary])
-        for _idx, _host_port in enumerate(_secondaries):
-            _host, _port = _split_host_port(_host_port)
-            # Enumerate doesn't accept a start argument until python 2.6.
-            setattr(sys.modules[__name__], "host%d" % (_idx + 1,), _host)
-            setattr(sys.modules[__name__], "port%d" % (_idx + 1,), _port)
+    # Try to find the primary for 10 seconds.
+    for _ in range(10):
+        _ismaster = pymongo.MongoClient(host, port).admin.command('ismaster')
+        # Master slave topologies don't use setName. The members also aren't
+        # aware of one another (e.g. ismaster doesn't have a 'hosts' field).
+        if 'setName' in _ismaster:
+            _primary = _ismaster.get('primary')
+            # No primary (yet?)
+            if _primary is None:
+                time.sleep(1)
+                continue
+            host, port = _split_host_port(_primary)
+            _secondaries = set(_ismaster['hosts']) - set([_primary])
+            for _idx, _host_port in enumerate(_secondaries):
+                _host, _port = _split_host_port(_host_port)
+                # Enumerate doesn't accept a start argument until python 2.6.
+                setattr(sys.modules[__name__], "host%d" % (_idx + 1,), _host)
+                setattr(sys.modules[__name__], "port%d" % (_idx + 1,), _port)
+        break
 except ConnectionFailure:
     pass