Test refactoring: get_pool() and pools_from_rs_client().

ajdavis · ajdavis · commit c1f6fece26d4 · 2013-10-08T10:40:16.000-04:00
diff --git a/test/test_legacy_connections.py b/test/test_legacy_connections.py
@@ -28,6 +28,7 @@
 from pymongo.errors import ConfigurationError
 from test import host, port, pair
 from test.test_replica_set_client import TestReplicaSetClientBase
+from test.utils import get_pool
 
 
 class TestConnection(unittest.TestCase):
@@ -92,8 +93,7 @@ def test_replica_set_connection(self):
 
         # To preserve legacy ReplicaSetConnection's behavior, max_size should
         # be None. Pool should handle this without error.
-        rs_state = c._MongoReplicaSetClient__rs_state
-        pool = rs_state.primary_member.pool
+        pool = get_pool(c)
         self.assertEqual(None, pool.max_size)
         c.end_request()
 
@@ -104,7 +104,8 @@ def test_replica_set_connection(self):
         )._MongoReplicaSetClient__net_timeout)
 
         for network_timeout in 'foo', 0, -1:
-            self.assertRaises(ConfigurationError,
+            self.assertRaises(
+                ConfigurationError,
                 ReplicaSetConnection, pair, replicaSet=self.name,
                 network_timeout=network_timeout)
 
diff --git a/test/test_replica_set_client.py b/test/test_replica_set_client.py
@@ -50,7 +50,7 @@
 from test.utils import (
     delay, assertReadFrom, assertReadFromAll, read_from_which_host,
     remove_all_users, assertRaisesExactly, TestRequestMixin, one,
-    server_started_with_auth)
+    server_started_with_auth, pools_from_rs_client, get_pool)
 
 
 class TestReplicaSetClientAgainstStandalone(unittest.TestCase):
@@ -704,8 +704,8 @@ def test_socket_error_marks_member_down(self):
         previous_writer = c._MongoReplicaSetClient__rs_state.writer
 
         def kill_sockets():
-            for member in c._MongoReplicaSetClient__rs_state.members:
-                for socket_info in member.pool.sockets:
+            for pool in pools_from_rs_client(c):
+                for socket_info in pool.sockets:
                     socket_info.sock.close()
 
         kill_sockets()
@@ -928,7 +928,7 @@ def test_operation_failure_without_request(self):
         # Ensure MongoReplicaSetClient doesn't close socket after it gets an
         # error response to getLastError. PYTHON-395.
         c = self._get_client(auto_start_request=False)
-        pool = c._MongoReplicaSetClient__rs_state.get(self.primary).pool
+        pool = get_pool(c)
         self.assertEqual(1, len(pool.sockets))
         old_sock_info = iter(pool.sockets).next()
         c.pymongo_test.test.drop()
@@ -948,7 +948,7 @@ def test_operation_failure_with_request(self):
         # error response to getLastError. PYTHON-395.
         c = self._get_client(auto_start_request=True)
         c.pymongo_test.test.find_one()
-        pool = c._MongoReplicaSetClient__rs_state.get(self.primary).pool
+        pool = get_pool(c)
 
         # Client reserved a socket for this thread
         self.assertTrue(isinstance(pool._get_request_state(), SocketInfo))
@@ -973,13 +973,10 @@ def test_auto_start_request(self):
 
         client = self._get_client(auto_start_request=True)
         self.assertTrue(client.auto_start_request)
-        pools = [member.pool for member in
-                 client._MongoReplicaSetClient__rs_state.members]
-
+        pools = pools_from_rs_client(client)
         self.assertInRequestAndSameSock(client, pools)
 
-        primary_pool = \
-            client._MongoReplicaSetClient__rs_state.get(client.primary).pool
+        primary_pool = get_pool(client)
 
         # Trigger the RSC to actually start a request on primary pool
         client.pymongo_test.test.find_one()
@@ -1008,9 +1005,7 @@ def test_auto_start_request(self):
         client.close()
 
         client = self._get_client()
-        pools = [mongo.pool for mongo in
-                 client._MongoReplicaSetClient__rs_state.members]
-
+        pools = pools_from_rs_client(client)
         self.assertNotInRequestAndDifferentSock(client, pools)
         client.start_request()
         self.assertInRequestAndSameSock(client, pools)
@@ -1021,8 +1016,7 @@ def test_auto_start_request(self):
     def test_nested_request(self):
         client = self._get_client(auto_start_request=True)
         try:
-            pools = [member.pool for member in
-                     client._MongoReplicaSetClient__rs_state.members]
+            pools = pools_from_rs_client(client)
             self.assertTrue(client.in_request())
 
             # Start and end request - we're still in "outer" original request
@@ -1064,8 +1058,7 @@ def test_nested_request(self):
     def test_request_threads(self):
         client = self._get_client()
         try:
-            pools = [member.pool for member in
-                     client._MongoReplicaSetClient__rs_state.members]
+            pools = pools_from_rs_client(client)
             self.assertNotInRequestAndDifferentSock(client, pools)
 
             started_request, ended_request = threading.Event(), threading.Event()
diff --git a/test/test_threads.py b/test/test_threads.py
@@ -23,22 +23,11 @@
 from test.utils import (joinall, remove_all_users,
                         server_started_with_auth, RendezvousThread)
 from test.test_client import get_client
-from pymongo.mongo_client import MongoClient
-from pymongo.replica_set_connection import MongoReplicaSetClient
+from test.utils import get_pool
 from pymongo.pool import SocketInfo, _closed
 from pymongo.errors import AutoReconnect, OperationFailure
 
 
-def get_pool(client):
-    if isinstance(client, MongoClient):
-        return client._MongoClient__pool
-    elif isinstance(client, MongoReplicaSetClient):
-        rs_state = client._MongoReplicaSetClient__rs_state
-        return rs_state[rs_state.writer].pool
-    else:
-        raise TypeError(str(client))
-
-
 class AutoAuthenticateThreads(threading.Thread):
 
     def __init__(self, collection, num):
diff --git a/test/utils.py b/test/utils.py
@@ -270,6 +270,22 @@ def assertReadFromAll(testcase, rsc, members, *args, **kwargs):
 
     testcase.assertEqual(members, used)
 
+def get_pool(client):
+    if isinstance(client, MongoClient):
+        return client._MongoClient__pool
+    elif isinstance(client, MongoReplicaSetClient):
+        rs_state = client._MongoReplicaSetClient__rs_state
+        return rs_state.primary_member.pool
+    else:
+        raise TypeError(str(client))
+
+def pools_from_rs_client(client):
+    """Get Pool instances from a MongoReplicaSetClient or ReplicaSetConnection.
+    """
+    return [
+        member.pool for member in
+        client._MongoReplicaSetClient__rs_state.members]
+
 class TestRequestMixin(object):
     """Inherit from this class and from unittest.TestCase to get some
     convenient methods for testing connection pools and requests
