Merge pull request #1151 from RussellLuo/feature-support-mocking

Add support for mocking MongoEngine based on mongomock

AUTHORS

@@ -233,3 +233,4 @@ that much better:
  * Ashley Whetter (https://github.com/AWhetter)
  * Paul-Armand Verhaegen (https://github.com/paularmand)
  * Steven Rossiter (https://github.com/BeardedSteve)
+ * Luo Peng (https://github.com/RussellLuo)

docs/changelog.rst

@@ -5,6 +5,7 @@ Changelog
 Changes in 0.10.5
 =================
 - Fix for reloading of strict with special fields. #1156
+- Add support for mocking MongoEngine based on mongomock. #1151
 
 Changes in 0.10.4
 =================

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)

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')