From 23f07fde5e88c42c42f26ca48edd10456c7babbf Mon Sep 17 00:00:00 2001
From: RussellLuo <luopeng.he@gmail.com>
Date: Sun, 8 Nov 2015 22:11:47 +0800
Subject: [PATCH] Add support for mocking MongoEngine based on mongomock

Using `mongomock://` scheme in URI enables the mocking. Fix #1045.
---
 mongoengine/connection.py | 21 ++++++++++++++++++---
 tests/test_connection.py  | 19 ++++++++++++++++++-
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/mongoengine/connection.py b/mongoengine/connection.py
index 4abca1ab..2ea572df 100644
--- a/mongoengine/connection.py
+++ b/mongoengine/connection.py
@@ -54,8 +54,11 @@ def register_connection(alias, name=None, host=None, port=None,
     }
 
     # Handle uri style connections
-    if "://" in conn_settings['host']:
-        uri_dict = uri_parser.parse_uri(conn_settings['host'])
+    conn_host = conn_settings['host']
+    if conn_host.startswith('mongomock://'):
+        conn_settings['is_mock'] = True
+    elif '://' in conn_host:
+        uri_dict = uri_parser.parse_uri(conn_host)
         conn_settings.update({
             'name': uri_dict.get('database') or name,
             'username': uri_dict.get('username'),
@@ -106,7 +109,19 @@ def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
         conn_settings.pop('password', None)
         conn_settings.pop('authentication_source', None)
 
-        connection_class = MongoClient
+        is_mock = conn_settings.pop('is_mock', None)
+        if is_mock:
+            # Use MongoClient from mongomock
+            try:
+                import mongomock
+            except ImportError:
+                raise RuntimeError('You need mongomock installed '
+                                   'to mock MongoEngine.')
+            connection_class = mongomock.MongoClient
+        else:
+            # Use MongoClient from pymongo
+            connection_class = MongoClient
+
         if 'replicaSet' in conn_settings:
             # Discard port since it can't be used on MongoReplicaSetClient
             conn_settings.pop('port', None)
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 1b7b7a22..4466ee73 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -8,6 +8,7 @@ try:
     import unittest2 as unittest
 except ImportError:
     import unittest
+from nose.plugins.skip import SkipTest
 
 import pymongo
 from bson.tz_util import utc
@@ -51,6 +52,22 @@ class ConnectionTest(unittest.TestCase):
         conn = get_connection('testdb')
         self.assertTrue(isinstance(conn, pymongo.mongo_client.MongoClient))
 
+    def test_connect_in_mocking(self):
+        """Ensure that the connect() method works properly in mocking.
+        """
+        try:
+            import mongomock
+        except ImportError:
+            raise SkipTest('you need mongomock installed to run this testcase')
+
+        connect('mongoenginetest', host='mongomock://localhost')
+        conn = get_connection()
+        self.assertTrue(isinstance(conn, mongomock.MongoClient))
+
+        connect('mongoenginetest2', host='mongomock://localhost', alias='testdb')
+        conn = get_connection('testdb')
+        self.assertTrue(isinstance(conn, mongomock.MongoClient))
+
     def test_disconnect(self):
         """Ensure that the disconnect() method works properly
         """
@@ -151,7 +168,7 @@ class ConnectionTest(unittest.TestCase):
             self.assertRaises(ConnectionError, get_db, 'test1')
 
         # Authentication succeeds with "authSource"
-        test_conn2 = connect(
+        connect(
             'mongoenginetest', alias='test2',
             host=('mongodb://username2:password@localhost/'
                   'mongoenginetest?authSource=admin')