From 2c631faa6cb0ff925be380523ac5b54d9713efd2 Mon Sep 17 00:00:00 2001
From: Shane Harvey <shnhrv@gmail.com>
Date: Fri, 8 May 2020 15:17:11 -0700
Subject: [PATCH] PYTHON-2239 Avoid 30 second stalls in TestMongoClientFailover

---
 test/pymongo_mocks.py |  4 ++--
 test/test_client.py   | 22 ++++++++++++++++++----
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/test/pymongo_mocks.py b/test/pymongo_mocks.py
index 7ed2e8a30..bd3e1ae22 100644
--- a/test/pymongo_mocks.py
+++ b/test/pymongo_mocks.py
@@ -51,7 +51,7 @@ class MockPool(Pool):
             + 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 @@ class MockMonitor(Monitor):
             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)
diff --git a/test/test_client.py b/test/test_client.py
index 0540786c5..ef1b19485 100644
--- a/test/test_client.py
+++ b/test/test_client.py
@@ -48,6 +48,7 @@ from pymongo.errors import (AutoReconnect,
                             InvalidURI,
                             NetworkTimeout,
                             OperationFailure,
+                            ServerSelectionTimeoutError,
                             WriteConcernError)
 from pymongo.monitoring import (ServerHeartbeatListener,
                                 ServerHeartbeatStartedEvent)
@@ -1731,6 +1732,7 @@ class TestMongoClientFailover(MockClientTest):
                 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 @@ class TestMongoClientFailover(MockClientTest):
             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 @@ class TestMongoClientFailover(MockClientTest):
         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 @@ class TestMongoClientFailover(MockClientTest):
                 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 @@ class TestMongoClientFailover(MockClientTest):
 
             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.