Fix connecting to a list of hosts

mongoengine/connection.py

@@ -1,5 +1,5 @@
 from pymongo import MongoClient, ReadPreference, uri_parser
-from mongoengine.python_support import IS_PYMONGO_3
+from mongoengine.python_support import (IS_PYMONGO_3, str_types)
 
 __all__ = ['ConnectionError', 'connect', 'register_connection',
            'DEFAULT_CONNECTION_NAME']
@@ -56,14 +56,21 @@ def register_connection(alias, name=None, host=None, port=None,
         'authentication_source': authentication_source
     }
 
-    # Handle uri style connections
     conn_host = conn_settings['host']
-    if conn_host.startswith('mongomock://'):
+    # host can be a list or a string, so if string, force to a list
+    if isinstance(conn_host, str_types):
+        conn_host = [conn_host]
+
+    resolved_hosts = []
+    for entity in conn_host:
+        # Handle uri style connections
+        if entity.startswith('mongomock://'):
             conn_settings['is_mock'] = True
             # `mongomock://` is not a valid url prefix and must be replaced by `mongodb://`
-        conn_settings['host'] = conn_host.replace('mongomock://', 'mongodb://', 1)
-    elif '://' in conn_host:
-        uri_dict = uri_parser.parse_uri(conn_host)
+            resolved_hosts.append(entity.replace('mongomock://', 'mongodb://', 1))
+        elif '://' in entity:
+            uri_dict = uri_parser.parse_uri(entity)
+            resolved_hosts.append(entity)
             conn_settings.update({
                 'name': uri_dict.get('database') or name,
                 'username': uri_dict.get('username'),
@@ -75,6 +82,9 @@ def register_connection(alias, name=None, host=None, port=None,
                 conn_settings['replicaSet'] = True
             if 'authsource' in uri_options:
                 conn_settings['authentication_source'] = uri_options['authsource']
+        else:
+            resolved_hosts.append(entity)
+    conn_settings['host'] = resolved_hosts
 
     # Deprecated parameters that should not be passed on
     kwargs.pop('slaves', None)

tests/test_connection.py

@@ -88,6 +88,40 @@ class ConnectionTest(unittest.TestCase):
         conn = get_connection('testdb7')
         self.assertTrue(isinstance(conn, mongomock.MongoClient))
 
+    def test_connect_with_host_list(self):
+        """Ensure that the connect() method works when host is a list
+
+        Uses mongomock to test w/o needing multiple mongod/mongos processes
+        """
+        try:
+            import mongomock
+        except ImportError:
+            raise SkipTest('you need mongomock installed to run this testcase')
+
+        connect(host=['mongomock://localhost'])
+        conn = get_connection()
+        self.assertTrue(isinstance(conn, mongomock.MongoClient))
+
+        connect(host=['mongodb://localhost'], is_mock=True,  alias='testdb2')
+        conn = get_connection('testdb2')
+        self.assertTrue(isinstance(conn, mongomock.MongoClient))
+
+        connect(host=['localhost'], is_mock=True,  alias='testdb3')
+        conn = get_connection('testdb3')
+        self.assertTrue(isinstance(conn, mongomock.MongoClient))
+
+        connect(host=['mongomock://localhost:27017', 'mongomock://localhost:27018'], alias='testdb4')
+        conn = get_connection('testdb4')
+        self.assertTrue(isinstance(conn, mongomock.MongoClient))
+
+        connect(host=['mongodb://localhost:27017', 'mongodb://localhost:27018'], is_mock=True,  alias='testdb5')
+        conn = get_connection('testdb5')
+        self.assertTrue(isinstance(conn, mongomock.MongoClient))
+
+        connect(host=['localhost:27017', 'localhost:27018'], is_mock=True,  alias='testdb6')
+        conn = get_connection('testdb6')
+        self.assertTrue(isinstance(conn, mongomock.MongoClient))
+
     def test_disconnect(self):
         """Ensure that the disconnect() method works properly
         """