PYTHON-2744 Run LB tests against non-LB clusters (#638)

Fix serviceId fallback to make spec test pass.
Fix socket leak when SocketInfo connection handshake fails.

(cherry picked from commit bf78a9b)

Author: ShaneHarvey

pymongo/pool.py

@@ -264,6 +264,9 @@ def _set_keepalive_times(sock):
 # main thread, to avoid the deadlock. See PYTHON-607.
 u'foo'.encode('idna')
 
+# Remove after PYTHON-2712
+_MOCK_SERVICE_ID = False
+
 
 def _raise_connection_failure(address, error, msg_prefix=None):
     """Convert a socket.error to ConnectionFailure and raise it."""
@@ -618,7 +621,7 @@ def _ismaster(self, cluster_time, topology_version,
         doc = self.command('admin', cmd, publish_events=False,
                            exhaust_allowed=awaitable)
         # PYTHON-2712 will remove this topologyVersion fallback logic.
-        if self.opts.load_balanced:
+        if self.opts.load_balanced and _MOCK_SERVICE_ID:
             process_id = doc.get('topologyVersion', {}).get('processId')
             doc.setdefault('serviceId', process_id)
         ismaster = IsMaster(doc, awaitable=awaitable)
@@ -645,7 +648,7 @@ def _ismaster(self, cluster_time, topology_version,
         if self.opts.load_balanced:
             if not ismaster.service_id:
                 raise ConfigurationError(
-                    'Driver attempted to initialize in load balancing mode'
+                    'Driver attempted to initialize in load balancing mode,'
                     ' but the server does not support this mode')
             self.service_id = ismaster.service_id
             self.generation = self.pool_gen.get(self.service_id)
@@ -1311,11 +1314,11 @@ def connect(self, all_credentials=None):
             raise
 
         sock_info = SocketInfo(sock, self, self.address, conn_id)
-        if self.handshake:
-            sock_info.ismaster(all_credentials)
-            self.is_writable = sock_info.is_writable
-
         try:
+            if self.handshake:
+                sock_info.ismaster(all_credentials)
+                self.is_writable = sock_info.is_writable
+
             sock_info.check_auth(all_credentials)
         except BaseException:
             sock_info.close_socket(ConnectionClosedReason.ERROR)

test/__init__.py

@@ -97,6 +97,9 @@
 SINGLE_MONGOS_LB_URI = os.environ.get("SINGLE_MONGOS_LB_URI")
 MULTI_MONGOS_LB_URI = os.environ.get("MULTI_MONGOS_LB_URI")
 if TEST_LOADBALANCER:
+    # Remove after PYTHON-2712
+    from pymongo import pool
+    pool._MOCK_SERVICE_ID = True
     res = parse_uri(SINGLE_MONGOS_LB_URI)
     host, port = res['nodelist'][0]
     db_user = res['username'] or db_user

test/load_balancer/test_crud_unified.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 import sys
-import unittest
 
 sys.path[0:0] = [""]
 
+from test import unittest
 from test.test_crud_unified import *
 
 if __name__ == '__main__':

test/load_balancer/test_dns.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 import sys
-import unittest
 
 sys.path[0:0] = [""]
 
+from test import unittest
 from test.test_dns import *
 
 if __name__ == '__main__':

test/load_balancer/test_load_balancer.py

@@ -14,38 +14,12 @@
 
 """Test the Load Balancer unified spec tests."""
 
-import os
 import sys
 
 sys.path[0:0] = [""]
 
-from test import unittest, IntegrationTest, client_context
-from test.utils import get_pool
-from test.unified_format import generate_test_classes
-
-# Location of JSON test specifications.
-TEST_PATH = os.path.join(
-    os.path.dirname(os.path.realpath(__file__)), 'unified')
-
-# Generate unified tests.
-globals().update(generate_test_classes(TEST_PATH, module=__name__))
-
-
-class TestLB(IntegrationTest):
-    @client_context.require_load_balancer
-    def test_unpin_committed_transaction(self):
-        pool = get_pool(self.client)
-        with self.client.start_session() as session:
-            with session.start_transaction():
-                self.assertEqual(pool.active_sockets, 0)
-                self.db.test.insert_one({}, session=session)
-                self.assertEqual(pool.active_sockets, 1)  # Pinned.
-            self.assertEqual(pool.active_sockets, 1)  # Still pinned.
-        self.assertEqual(pool.active_sockets, 0)  # Unpinned.
-
-    def test_client_can_be_reopened(self):
-        self.client.close()
-        self.db.test.find_one({})
+from test import unittest
+from test.test_load_balancer import *
 
 
 if __name__ == "__main__":

test/load_balancer/test_retryable_change_stream.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 import sys
-import unittest
 
 sys.path[0:0] = [""]
 
+from test import unittest
 from test.test_change_stream import *
 
 if __name__ == '__main__':

test/load_balancer/test_retryable_reads.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 import sys
-import unittest
 
 sys.path[0:0] = [""]
 
+from test import unittest
 from test.test_retryable_reads import *
 
 if __name__ == '__main__':

test/load_balancer/test_retryable_writes.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 import sys
-import unittest
 
 sys.path[0:0] = [""]
 
+from test import unittest
 from test.test_retryable_writes import *
 
 if __name__ == '__main__':

test/load_balancer/test_transactions_unified.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 import sys
-import unittest
 
 sys.path[0:0] = [""]
 
+from test import unittest
 from test.test_transactions_unified import *
 
 if __name__ == '__main__':

test/load_balancer/test_uri_options.py

@@ -13,10 +13,10 @@
 # limitations under the License.
 
 import sys
-import unittest
 
 sys.path[0:0] = [""]
 
+from test import unittest
 from test.test_uri_spec import *
 
 if __name__ == '__main__':