Fix test_ship_of_theseus.

Author: ajdavis

test/high_availability/ha_tools.py

@@ -269,10 +269,18 @@ def create_sharded_cluster(num_routers=3):
 
 # Connect to a random member
 def get_client():
-    return pymongo.MongoClient(
-        nodes.keys(),
-        read_preference=ReadPreference.PRIMARY_PREFERRED,
-        use_greenlets=use_greenlets)
+    # Attempt a direct connection to each node until one succeeds. Using a
+    # non-PRIMARY read preference allows us to use the node even if it's a
+    # secondary.
+    for i, node in enumerate(nodes.keys()):
+        try:
+            return pymongo.MongoClient(
+                node,
+                read_preference=ReadPreference.PRIMARY_PREFERRED,
+                use_greenlets=use_greenlets)
+        except pymongo.errors.ConnectionFailure:
+            if i == len(nodes.keys()) - 1:
+                raise
 
 
 def get_mongos_seed_list():
@@ -301,7 +309,7 @@ def get_primary():
         assert len(primaries) <= 1
         if primaries:
             return primaries[0]
-    except pymongo.errors.ConnectionFailure:
+    except (pymongo.errors.ConnectionFailure, pymongo.errors.OperationFailure):
         pass
 
     return None
@@ -379,6 +387,7 @@ def add_member(auth=False):
     global cur_port
     host = '%s:%d' % (hostname, cur_port)
     primary = get_primary()
+    assert primary
     c = pymongo.MongoClient(primary, use_greenlets=use_greenlets)
     config = c.local.system.replset.find_one()
     _id = max([member['_id'] for member in config['members']]) + 1

test/high_availability/test_ha.py

@@ -1082,7 +1082,21 @@ def test_ship_of_theseus(self):
         primary = ha_tools.get_primary()
         secondary1 = ha_tools.get_random_secondary()
 
-        new_hosts = [ha_tools.add_member() for _ in range(3)]
+        new_hosts = []
+        for i in range(3):
+            new_hosts.append(ha_tools.add_member())
+
+            # RS closes all connections after reconfig.
+            for j in xrange(30):
+                try:
+                    if ha_tools.get_primary():
+                        break
+                except (ConnectionFailure, OperationFailure):
+                    pass
+
+                sleep(1)
+            else:
+                self.fail("Couldn't recover from reconfig")
 
         # Wait for new members to join.
         for _ in xrange(120):
@@ -1094,6 +1108,7 @@ def test_ship_of_theseus(self):
             self.fail("New secondaries didn't join")
 
         ha_tools.kill_members([primary, secondary1], 9)
+        sleep(5)
 
         # Wait for primary.
         for _ in xrange(30):
@@ -1132,6 +1147,18 @@ def test_ship_of_theseus(self):
 
         ha_tools.restart_members([primary, secondary1])
 
+        # Wait for members to figure out they're secondaries.
+        for _ in xrange(30):
+            try:
+                if len(ha_tools.get_secondaries()) == 2:
+                    break
+            except ConnectionFailure:
+                pass
+
+            sleep(1)
+        else:
+            self.fail("Original members didn't become secondaries")
+
         # Should be able to reconnect to set again.
         sleep(2 * MONITOR_INTERVAL)
         find_one(read_preference=SECONDARY)