Apply multithreading unittests to new ReplicaSetConnection class.

Author: A. Jesse Jiryu Davis

test/test_replica_set_connection.py

@@ -42,8 +42,7 @@
 port = os.environ.get("DB_PORT", 27017)
 pair = '%s:%d' % (host, port)
 
-class TestConnection(unittest.TestCase):
-
+class TestConnectionReplicaSetBase(unittest.TestCase):
     def setUp(self):
         conn = Connection(pair)
         response = conn.admin.command('ismaster')
@@ -58,11 +57,12 @@ def setUp(self):
         else:
             raise SkipTest()
 
-    def _get_connection(self, **kwargs):
-        return ReplicaSetConnection(pair,
-                                    replicaSet=self.name,
-                                    **kwargs)
+def _get_connection(self, **kwargs):
+    return ReplicaSetConnection(pair,
+        replicaSet=self.name,
+        **kwargs)
 
+class TestConnection(TestConnectionReplicaSetBase):
     def test_connect(self):
         self.assertRaises(ConnectionFailure, ReplicaSetConnection,
                           "somedomainthatdoesntexist.org:27017",

test/test_threads.py

@@ -120,7 +120,16 @@ def run(self):
 class TestThreads(unittest.TestCase):
 
     def setUp(self):
-        self.db = get_connection().pymongo_test
+        self.db = self._get_connection().pymongo_test
+
+    def _get_connection(self):
+        """
+        Intended for overriding in TestThreadsReplicaSet. This method
+        returns a Connection here, and a ReplicaSetConnection in
+        test_threads_replica_set_connection.py.
+        """
+        # Regular test connection
+        return get_connection()
 
     def test_threading(self):
         self.db.drop_collection("test")
@@ -197,11 +206,17 @@ def test_low_network_timeout(self):
 
 
 class TestThreadsAuth(unittest.TestCase):
+    def _get_connection(self):
+        """
+        Intended for overriding in TestThreadsAuthReplicaSet. This method
+        returns a Connection here, and a ReplicaSetConnection in
+        test_threads_replica_set_connection.py.
+        """
+        # Regular test connection
+        return get_connection()
 
     def setUp(self):
-        self.conn = get_connection()
-
-        # Setup auth users
+        self.conn = self._get_connection()
         self.conn.admin.system.users.remove({})
         self.conn.admin.add_user('admin-user', 'password')
         try:
@@ -225,11 +240,11 @@ def tearDown(self):
         self.conn.drop_database('auth_test')
 
     def test_auto_auth_login(self):
-        conn = get_connection()
+        conn = self._get_connection()
         self.assertRaises(OperationFailure, conn.auth_test.test.find_one)
 
         # Admin auth
-        conn = get_connection()
+        conn = self._get_connection()
         conn.admin.authenticate("admin-user", "password")
 
         threads = []
@@ -242,7 +257,7 @@ def test_auto_auth_login(self):
             self.assertTrue(t.success)
 
         # Database-specific auth
-        conn = get_connection()
+        conn = self._get_connection()
         conn.auth_test.authenticate("test-user", "password")
 
         threads = []

test/test_threads_replica_set_connection.py

@@ -0,0 +1,116 @@
+# Copyright 2009-2010 10gen, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Test that pymongo is thread safe."""
+
+import unittest
+import os
+
+from nose.plugins.skip import SkipTest
+from pymongo.replica_set_connection import ReplicaSetConnection
+from test.test_threads import IgnoreAutoReconnect
+from test.test_threads import TestThreads as _testThreads
+from test.test_threads import TestThreadsAuth as _testThreadsAuth
+
+from test.test_replica_set_connection import (TestConnectionReplicaSetBase,
+                                              pair)
+
+from pymongo.errors import AutoReconnect
+
+class TestThreadsReplicaSet(TestConnectionReplicaSetBase, _testThreads):
+    def setUp(self):
+        """
+        Prepare to test all the same things that TestThreads tests, but do it
+        with a replica-set connection
+        """
+        TestConnectionReplicaSetBase.setUp(self)
+        _testThreads.setUp(self)
+
+    def _get_connection(self):
+        """
+        Override TestThreads, so its tests run on a ReplicaSetConnection
+        instead of a regular Connection.
+        """
+        return ReplicaSetConnection(pair, replicaSet=self.name)
+
+    def test_low_network_timeout(self):
+        """
+        Override the similar test as TestThreads.test_low_connect_timeout(),
+        but use the newer connectTimeoutMS and socketTimeoutMS arguments instead
+        of the soon-to-be-deprecated network_timeout.
+        """
+        db = None
+        n = 10
+
+        # First test connection timeout
+        i = 0
+        while db is None and i < n:
+            try:
+                db = ReplicaSetConnection(
+                    pair,
+                    replicaSet=self.name,
+                    connectTimeoutMS=1
+                ).pymongo_test
+            except AutoReconnect:
+                i += 1
+        if i == n:
+            raise SkipTest()
+
+        # ... then test socket timeout
+        i = 0
+        while db is None and i < n:
+            try:
+                db = ReplicaSetConnection(
+                    pair,
+                    replicaSet=self.name,
+                    socketTimeoutMS=1
+                ).pymongo_test
+            except AutoReconnect:
+                i += 1
+        if i == n:
+            raise SkipTest()
+
+        threads = []
+        for _ in range(4):
+            t = IgnoreAutoReconnect(db.test, 100)
+            t.start()
+            threads.append(t)
+
+        for t in threads:
+            t.join()
+
+
+class TestThreadsAuthReplicaSet(TestConnectionReplicaSetBase, _testThreadsAuth):
+
+    def setUp(self):
+        """
+        Prepare to test all the same things that TestThreads tests, but do it
+        with a replica-set connection
+        """
+        TestConnectionReplicaSetBase.setUp(self)
+        _testThreadsAuth.setUp(self)
+
+    def _get_connection(self):
+        """
+        Override TestThreadsAuth, so its tests run on a ReplicaSetConnection
+        instead of a regular Connection.
+        """
+        return ReplicaSetConnection(pair, replicaSet=self.name)
+
+if __name__ == "__main__":
+    suite = unittest.TestSuite([
+        unittest.makeSuite(TestThreadsReplicaSet),
+        unittest.makeSuite(TestThreadsAuthReplicaSet)
+    ])
+    unittest.TextTestRunner(verbosity=2).run(suite)