PYTHON-2239 Avoid 30 second stalls in TestMongoClientFailover

Author: ShaneHarvey

test/pymongo_mocks.py

@@ -51,7 +51,7 @@ def get_socket(self, all_credentials, checkout=False):
             + client.mock_members
             + client.mock_mongoses), "bad host: %s" % host_and_port
 
-        with Pool.get_socket(self, all_credentials) as sock_info:
+        with Pool.get_socket(self, all_credentials, checkout) as sock_info:
             sock_info.mock_host = self.mock_host
             sock_info.mock_port = self.mock_port
             yield sock_info
@@ -74,7 +74,7 @@ def __init__(
             pool,
             topology_settings)
 
-    def _check_once(self, metadata=None, cluster_time=None):
+    def _check_once(self):
         address = self._server_description.address
         response, rtt = self.client.mock_is_master('%s:%d' % address)
         return ServerDescription(address, IsMaster(response), rtt)

test/test_client.py

@@ -48,6 +48,7 @@
                             InvalidURI,
                             NetworkTimeout,
                             OperationFailure,
+                            ServerSelectionTimeoutError,
                             WriteConcernError)
 from pymongo.monitoring import (ServerHeartbeatListener,
                                 ServerHeartbeatStartedEvent)
@@ -1731,6 +1732,7 @@ def test_discover_primary(self):
                 mongoses=[],
                 host='b:2',  # Pass a secondary.
                 replicaSet='rs')
+            self.addCleanup(c.close)
 
             wait_until(lambda: len(c.nodes) == 3, 'connect')
             self.assertEqual(c.address, ('a', 1))
@@ -1760,7 +1762,10 @@ def test_reconnect(self):
             mongoses=[],
             host='b:2',  # Pass a secondary.
             replicaSet='rs',
-            retryReads=False)
+            retryReads=False,
+            serverSelectionTimeoutMS=100,
+        )
+        self.addCleanup(c.close)
 
         wait_until(lambda: len(c.nodes) == 3, 'connect')
 
@@ -1769,8 +1774,13 @@ def test_reconnect(self):
         c.kill_host('b:2')
         c.kill_host('c:3')
 
-        # MongoClient discovers it's alone.
+        # MongoClient discovers it's alone. The first attempt raises either
+        # ServerSelectionTimeoutError or AutoReconnect (from
+        # MockPool.get_socket).
         self.assertRaises(AutoReconnect, c.db.collection.find_one)
+        # The second attempt always raises ServerSelectionTimeoutError.
+        self.assertRaises(ServerSelectionTimeoutError,
+                          c.db.collection.find_one)
 
         # But it can reconnect.
         c.revive_host('a:1')
@@ -1789,7 +1799,9 @@ def _test_network_error(self, operation_callback):
                 host='a:1',
                 replicaSet='rs',
                 connect=False,
-                retryReads=False)
+                retryReads=False,
+                serverSelectionTimeoutMS=100)
+            self.addCleanup(c.close)
 
             # Set host-specific information so we can test whether it is reset.
             c.set_wire_version_range('a:1', 2, 6)
@@ -1799,7 +1811,9 @@ def _test_network_error(self, operation_callback):
 
             c.kill_host('a:1')
 
-            # MongoClient is disconnected from the primary.
+            # MongoClient is disconnected from the primary. This raises either
+            # ServerSelectionTimeoutError or AutoReconnect (from
+            # MockPool.get_socket).
             self.assertRaises(AutoReconnect, operation_callback, c)
 
             # The primary's description is reset.