modifying test about connection

mongoengine/connection.py

@@ -92,17 +92,21 @@ def register_connection(alias, name=None, host=None, port=None,
             if uri_dict.get('database'):
                 conn_settings['name'] = uri_dict.get('database')
 
-            for param in ('read_preference', 'username', 'password'):
+            for param in ('username', 'password'):
                 if uri_dict.get(param):
                     conn_settings[param] = uri_dict[param]
 
             uri_options = uri_dict['options']
-            if 'replicaset' in uri_options:
-                conn_settings['replicaSet'] = uri_options['replicaset']
+            if 'replicaset' in uri_options \
+                    or 'replicaset' in [kw_key.lower() for kw_key in kwargs.keys()]:
+                conn_settings['isreplicaSet'] = True
             if 'authsource' in uri_options:
                 conn_settings['authentication_source'] = uri_options['authsource']
             if 'authmechanism' in uri_options:
                 conn_settings['authentication_mechanism'] = uri_options['authmechanism']
+            if IS_PYMONGO_3 and 'readpreference' in uri_options \
+                    and 'readpreferencetags' in uri_options:
+                del conn_settings['read_preference']
         else:
             resolved_hosts.append(entity)
     conn_settings['host'] = resolved_hosts
@@ -176,19 +180,18 @@ def _clean_settings(settings_dict):
         # For replica set connections with PyMongo 2.x, use
         # MongoReplicaSetClient.
         # TODO remove this once we stop supporting PyMongo 2.x.
-        if 'replicaSet' in conn_settings and not IS_PYMONGO_3:
-            connection_class = MongoReplicaSetClient
-            conn_settings['hosts_or_uri'] = conn_settings.pop('host', None)
-
-            # hosts_or_uri has to be a string, so if 'host' was provided
-            # as a list, join its parts and separate them by ','
-            if isinstance(conn_settings['hosts_or_uri'], list):
-                conn_settings['hosts_or_uri'] = ','.join(
-                    conn_settings['hosts_or_uri'])
-
-            # Discard port since it can't be used on MongoReplicaSetClient
-            conn_settings.pop('port', None)
-
+        if 'isreplicaSet' in conn_settings.keys():
+            if not IS_PYMONGO_3:
+                connection_class = MongoReplicaSetClient
+                conn_settings['hosts_or_uri'] = conn_settings.pop('host', None)
+                # hosts_or_uri has to be a string, so if 'host' was provided
+                # as a list, join its parts and separate them by ','
+                if isinstance(conn_settings['hosts_or_uri'], list):
+                    conn_settings['hosts_or_uri'] = ','.join(
+                        conn_settings['hosts_or_uri'])
+                # Discard port since it can't be used on MongoReplicaSetClient
+                conn_settings.pop('port', None)
+            del conn_settings['isreplicaSet']
     # Iterate over all of the connection settings and if a connection with
     # the same parameters is already established, use it instead of creating
     # a new one.

tests/test_connection.py → tests/connections/test_connection.py

@@ -211,6 +211,78 @@ def test_uri_without_credentials_doesnt_override_conn_settings(self):
         # behavior. If the MongoDB URI would override the credentials
         self.assertRaises(OperationFailure, get_db)
 
+    @unittest.skip('require authentication enabled and meet needed prerequisites')
+    def test_uri_overide_kwargs_with_right_authSource_not_pass(self):
+        """Authenticated during building connection"""
+        ###as prerequisite, require two users at the time when authentication is disabled###
+        #no_auth_conn_obj = connect(db='mongoenginetest')
+        #no_auth_conn_obj.admin.system.users.remove({})
+        #no_auth_conn_obj.admin.add_user('rjxu', 'rjxp', read_only=True)
+        #no_auth_conn_obj.mongoenginetest.add_user('rjxtu','rjxtp', read_only=True)
+
+        auth_conn_obj = connect(
+            db='mongoenginetest',
+            alias='authtest',
+            username='rjxu',
+            password='rjxp',
+            host="mongodb://rjxtu:rjxtp@localhost:27017" +
+                 "/?authSource=mongoenginetest",
+            authentication_source='admin',
+            authentication_mechanism='MONGODB-X509'
+        )
+
+        if not IS_PYMONGO_3:
+            #attempt to read data from a collection authored by admin
+            def get_data_case():
+                auth_db_obj = auth_conn_obj['admin']
+                coll_from_auth_db = auth_db_obj['system.users']
+                return [sys_user for sys_user in coll_from_auth_db.find()]
+            try:
+                get_data_case()
+            except:
+                self.assertRaises(OperationFailure, get_data_case)
+        else:
+            credentials_obj = auth_conn_obj._MongoClient__options.credentials
+            self.assertEqual(credentials_obj.username, 'rjxtu')
+            self.assertEqual(credentials_obj.source, 'mongoenginetest')
+            self.assertEqual(credentials_obj.mechanism, 'DEFAULT')
+
+    @unittest.skip('require authentication enabled and meet needed prerequisite')
+    def test_post_auth_mongo_client_with_get_db(self):
+        '''Post authentication at get db stage'''
+        ###as prerequisite, require one user at the time when authentication is disabled###
+        #no_auth_conn_obj = connect(db='mongoenginetest')
+        #no_auth_conn_obj.admin.system.users.remove({})
+        #no_auth_conn_obj.mongoenginetest.add_user('rjxtu', 'rjxtp', read_only=True)
+
+        no_auth_conn_obj = connect(
+            db='mongoenginetest',
+            username='rjxtu',
+            password='rjxtp',
+            host="mongodb://localhost:27017",
+            authentication_source='mongoenginetest',
+            authentication_mechanism='DEFAULT'
+        )
+
+        #read data from auth or non-auth conn
+        def get_data_case(is_get_db=False):
+            if not is_get_db:
+                _db_obj = no_auth_conn_obj['mongoenginetest']
+            else:
+                _db_obj = get_db()
+            coll_from_db = _db_obj['testcoll']
+            return [_td for _td in coll_from_db.find()]
+
+        #
+        try:
+            get_data_case()
+        except:
+            self.assertRaises(OperationFailure, get_data_case)
+        #
+        self.assertRaises(AssertionError, self.assertRaises,
+                          OperationFailure, get_data_case, True)
+
+
     def test_connect_uri_with_authsource(self):
         """Ensure that the connect() method works well with `authSource`
         option in the URI.
@@ -319,38 +391,6 @@ def test_write_concern(self):
             self.assertEqual(dict(conn1.write_concern), {'w': 1, 'j': True})
             self.assertEqual(dict(conn2.write_concern), {'w': 1, 'j': True})
 
-    def test_connect_with_replicaset_via_uri(self):
-        """Ensure connect() works when specifying a replicaSet via the
-        MongoDB URI.
-        """
-        if IS_PYMONGO_3:
-            c = connect(host='mongodb://localhost/test?replicaSet=local-rs')
-            db = get_db()
-            self.assertTrue(isinstance(db, pymongo.database.Database))
-            self.assertEqual(db.name, 'test')
-        else:
-            # PyMongo < v3.x raises an exception:
-            # "localhost:27017 is not a member of replica set local-rs"
-            with self.assertRaises(MongoEngineConnectionError):
-                c = connect(host='mongodb://localhost/test?replicaSet=local-rs')
-
-    def test_connect_with_replicaset_via_kwargs(self):
-        """Ensure connect() works when specifying a replicaSet via the
-        connection kwargs
-        """
-        if IS_PYMONGO_3:
-            c = connect(replicaset='local-rs')
-            self.assertEqual(c._MongoClient__options.replica_set_name,
-                             'local-rs')
-            db = get_db()
-            self.assertTrue(isinstance(db, pymongo.database.Database))
-            self.assertEqual(db.name, 'test')
-        else:
-            # PyMongo < v3.x raises an exception:
-            # "localhost:27017 is not a member of replica set local-rs"
-            with self.assertRaises(MongoEngineConnectionError):
-                c = connect(replicaset='local-rs')
-
     def test_datetime(self):
         connect('mongoenginetest', tz_aware=True)
         d = datetime.datetime(2010, 5, 5, tzinfo=utc)

tests/connections/test_replicaset_connection.py

@@ -0,0 +1,147 @@
+import unittest
+
+from pymongo import (ReadPreference,
+                     read_preferences)
+
+from mongoengine.python_support import IS_PYMONGO_3
+
+if IS_PYMONGO_3:
+    from pymongo import MongoClient
+    CONN_CLASS = MongoClient
+    READ_PREF = ReadPreference.SECONDARY
+else:
+    from pymongo import ReplicaSetConnection
+    CONN_CLASS = ReplicaSetConnection
+    READ_PREF = ReadPreference.SECONDARY_ONLY
+
+import pymongo
+import mongoengine
+from mongoengine import *
+from mongoengine.connection import (MongoEngineConnectionError,
+                                    get_db)
+
+
+class ConnectionTest(unittest.TestCase):
+
+    def setUp(self):
+        mongoengine.connection._connection_settings = {}
+        mongoengine.connection._connections = {}
+        mongoengine.connection._dbs = {}
+
+    def tearDown(self):
+        mongoengine.connection._connection_settings = {}
+        mongoengine.connection._connections = {}
+        mongoengine.connection._dbs = {}
+
+    def test_replicaset_uri_passes_read_preference(self):
+        """Requires a replica set called "rs" on port 27017
+        """
+
+        try:
+            conn = connect(db='mongoenginetest',
+                           host="mongodb://localhost/mongoenginetest?replicaSet=local-rs",
+                           read_preference=READ_PREF)
+        except MongoEngineConnectionError as e:
+            return
+
+        if not isinstance(conn, CONN_CLASS):
+            # really???
+            return
+
+        self.assertEqual(conn.read_preference, READ_PREF)
+
+    def test_connect_with_replicaset_via_kwargs(self):
+        """Ensure connect() works when specifying a replicaSet via the
+        connection kwargs
+        """
+        if IS_PYMONGO_3:
+            c = connect(replicaset='local-rs')
+            self.assertEqual(c._MongoClient__options.replica_set_name,
+                             'local-rs')
+            db = get_db()
+            self.assertTrue(isinstance(db, pymongo.database.Database))
+            self.assertEqual(db.name, 'test')
+        else:
+            # PyMongo < v3.x raises an exception:
+            # "localhost:27017 is not a member of replica set local-rs"
+            with self.assertRaises(MongoEngineConnectionError):
+                c = connect(replicaset='local-rs')
+
+    def test_connect_with_replicaset_via_uri(self):
+        """Ensure connect() works when specifying a replicaSet via the
+        MongoDB URI.
+        """
+        if IS_PYMONGO_3:
+            c = connect(host='mongodb://localhost/test?replicaSet=local-rs')
+            db = get_db()
+            self.assertTrue(isinstance(db, pymongo.database.Database))
+            self.assertEqual(db.name, 'test')
+        else:
+            # PyMongo < v3.x raises an exception:
+            # "localhost:27017 is not a member of replica set local-rs"
+            with self.assertRaises(MongoEngineConnectionError):
+                c = connect(host='mongodb://localhost/test?replicaSet=local-rs')
+
+    def test_read_preference_from_replica_set_in_uri_as_host(self):
+        '''read preference from replica set cluster'''
+        #test case about uri option overrides kwargs and read_preference with tag sets
+        if IS_PYMONGO_3:
+            conn_obj = connect(
+                db='testrjx',
+                host="mongodb://localhost:27017/?replicaset=dev_rs" +
+                     "&readpreference=nearest&readpreferencetags=",
+                read_preference=ReadPreference.SECONDARY
+            )
+            read_preference_obj = conn_obj.read_preference
+            self.assertEqual(read_preference_obj.mode, ReadPreference.NEAREST.mode)
+            self.assertEqual(read_preference_obj.tag_sets, [{}])
+        else:
+            if pymongo.version_tuple[1] < 9: # like v2.8
+                with self.assertRaises(MongoEngineConnectionError):
+                    c_obj = connect(
+                        db='testrjx',
+                        host="mongodb://localhost:27018/?replicaset=dev_rs"
+                        + "&read_preference=nearest&readpreferencetags=dc:east,use:dev",
+                        read_preference=ReadPreference.SECONDARY,
+                        readpreferencetags=[{}] #tag_sets as an alternative
+                    )
+                    self.assertEqual(c_obj.read_preference, ReadPreference.NEAREST)
+                    self.assertEqual(c_obj.tag_sets, [{'dc':'east','use':'dev'}])
+            else:
+                #for 2.9 as transition version, also accept read_preference object subclassing _ServerMode
+                with self.assertRaises(MongoEngineConnectionError):
+                    c_obj = connect(
+                        db='testrjx',
+                        host="mongodb://localhost:27018/?replicaset=dev_rs",
+                        read_preference=read_preferences.SecondaryPreferred(
+                            [{'dc':'east','use':'dev'}]
+                        ),
+                    )
+                    self.assertEqual(c_obj.read_preference, ReadPreference.SECONDARY_PREFERRED)
+                    self.assertEqual(c_obj.tag_sets, [{'dc':'east','use':'dev'}])
+
+    def test_replica_set_as_uri_or_kwargs(self):
+        '''replicaset as uri option or kwargs'''
+        # dev_rs as real replica set name, while test_rs is false as comparision
+        # replicaset case insensitive
+        if not IS_PYMONGO_3: #like 2.9 and 2.8
+            # uri opt overrides kwargs for connecting replica set
+            with self.assertRaises(MongoEngineConnectionError):
+                c_obj = connect(
+                    db='testrjx',
+                    host="mongodb://localhost:27018/?replicaset=dev_rs",
+                    replicaset='test_rs'
+                )
+                self.assertEqual(c_obj._MongoReplicaSetClient__name, 'dev_rs')
+        else: #like 3.3
+            #here, kwargs overrides uri option
+            #though wrong replica set to be connected, it is ok until issuing real data operation
+            c_obj = connect(
+                db='testrjx',
+                host="mongodb://localhost:27017/?replicaSet=dev_rs",
+                replicaSet='test_rs'
+            )
+            self.assertEqual(c_obj._MongoClient__options.replica_set_name, 'test_rs')
+
+if __name__ == '__main__':
+    unittest.main()