diff --git a/AUTHORS b/AUTHORS index 66e181c9..b087f41c 100644 --- a/AUTHORS +++ b/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) diff --git a/docs/changelog.rst b/docs/changelog.rst index c354be18..13b14a0f 100644 --- a/docs/changelog.rst +++ b/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 ================= 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')