From 13897db6d3a953077a3fc4c4d555141e959d42da Mon Sep 17 00:00:00 2001
From: Emmanuel Leblond <emmanuel.leblond@gmail.com>
Date: Wed, 20 Jan 2016 11:06:45 +0100
Subject: [PATCH] Fix mongomock url prefix error during connection

---
 mongoengine/connection.py |  5 +++++
 tests/test_connection.py  | 24 ++++++++++++++++++++++--
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/mongoengine/connection.py b/mongoengine/connection.py
index cde3ccd4..4055a9b6 100644
--- a/mongoengine/connection.py
+++ b/mongoengine/connection.py
@@ -38,8 +38,11 @@ def register_connection(alias, name=None, host=None, port=None,
     :param username: username to authenticate with
     :param password: password to authenticate with
     :param authentication_source: database to authenticate against
+    :param is_mock: explicitly use mongomock for this connection
+        (can also be done by using `mongomock://` as db host prefix)
     :param kwargs: allow ad-hoc parameters to be passed into the pymongo driver
 
+    .. versionchanged:: 0.10.6 - added mongomock support
     """
     global _connection_settings
 
@@ -57,6 +60,8 @@ def register_connection(alias, name=None, host=None, port=None,
     conn_host = conn_settings['host']
     if conn_host.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)
         conn_settings.update({
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 4466ee73..b2f7406e 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -64,8 +64,28 @@ class ConnectionTest(unittest.TestCase):
         conn = get_connection()
         self.assertTrue(isinstance(conn, mongomock.MongoClient))
 
-        connect('mongoenginetest2', host='mongomock://localhost', alias='testdb')
-        conn = get_connection('testdb')
+        connect('mongoenginetest2', host='mongomock://localhost', alias='testdb2')
+        conn = get_connection('testdb2')
+        self.assertTrue(isinstance(conn, mongomock.MongoClient))
+
+        connect('mongoenginetest3', host='mongodb://localhost', is_mock=True, alias='testdb3')
+        conn = get_connection('testdb3')
+        self.assertTrue(isinstance(conn, mongomock.MongoClient))
+
+        connect('mongoenginetest4', is_mock=True, alias='testdb4')
+        conn = get_connection('testdb4')
+        self.assertTrue(isinstance(conn, mongomock.MongoClient))
+
+        connect(host='mongodb://localhost:27017/mongoenginetest5', is_mock=True, alias='testdb5')
+        conn = get_connection('testdb5')
+        self.assertTrue(isinstance(conn, mongomock.MongoClient))
+
+        connect(host='mongomock://localhost:27017/mongoenginetest6', alias='testdb6')
+        conn = get_connection('testdb6')
+        self.assertTrue(isinstance(conn, mongomock.MongoClient))
+
+        connect(host='mongomock://localhost:27017/mongoenginetest7', is_mock=True, alias='testdb7')
+        conn = get_connection('testdb7')
         self.assertTrue(isinstance(conn, mongomock.MongoClient))
 
     def test_disconnect(self):