Merge pull request #304 from touilleMan/fix-test-and-mongomock

Fix test and mongomock

Author: touilleMan

AUTHORS

@@ -35,3 +35,4 @@ that much better:
 * Denny Huang - https://github.com/denny0223
 * Stefan Wojcik - https://github.com/wojcikstefan
 * John Cass - https://github.com/jcass77
+* Aly Sivji - https://github.com/alysivji

flask_mongoengine/connection.py

@@ -1,11 +1,16 @@
 import mongoengine
 from pymongo import ReadPreference, uri_parser
 
+
 __all__ = (
     'create_connections', 'get_connection_settings', 'InvalidSettingsError',
 )
 
 
+MONGODB_CONF_VARS = ('MONGODB_ALIAS', 'MONGODB_DB', 'MONGODB_HOST', 'MONGODB_IS_MOCK',
+                     'MONGODB_PASSWORD', 'MONGODB_PORT', 'MONGODB_USERNAME')
+
+
 class InvalidSettingsError(Exception):
     pass
 
@@ -24,7 +29,15 @@ def _sanitize_settings(settings):
 
     # Handle uri style connections
     if "://" in resolved_settings.get('host', ''):
-        uri_dict = uri_parser.parse_uri(resolved_settings['host'])
+        # this section pulls the database name from the URI
+        # PyMongo requires URI to start with mongodb:// to parse
+        # this workaround allows mongomock to work
+        uri_to_check = resolved_settings['host']
+
+        if uri_to_check.startswith('mongomock://'):
+            uri_to_check = uri_to_check.replace('mongomock://', 'mongodb://')
+
+        uri_dict = uri_parser.parse_uri(uri_to_check)
         resolved_settings['db'] = uri_dict['database']
 
     # Add a default name param or use the "db" key if exists
@@ -74,10 +87,10 @@ def get_connection_settings(config):
         else:
             return _sanitize_settings(settings)
 
-    # If "MONGODB_SETTINGS" doesn't exist, sanitize all the keys starting with
-    # "MONGODB_" as if they all describe a single connection.
+    # If "MONGODB_SETTINGS" doesn't exist, sanitize the "MONGODB_" keys
+    # as if they all describe a single connection.
     else:
-        config = dict((k, v) for k, v in config.items() if k.startswith('MONGODB_'))  # ugly dict comprehention in order to support python 2.6
+        config = dict((k, v) for k, v in config.items() if k in MONGODB_CONF_VARS)  # ugly dict comprehention in order to support python 2.6
         return _sanitize_settings(config)
 
 

tests/__init__.py

@@ -1,6 +1,6 @@
 import unittest
-
 import flask
+import mongoengine
 
 
 class FlaskMongoEngineTestCase(unittest.TestCase):
@@ -12,6 +12,13 @@ def setUp(self):
         self.app.config['TESTING'] = True
         self.ctx = self.app.app_context()
         self.ctx.push()
+        # Mongoengine keep a global state of the connections that must be
+        # reset before each test.
+        # Given it doesn't expose any method to get the list of registered
+        # connections, we have to do the cleaning by hand...
+        mongoengine.connection._connection_settings.clear()
+        mongoengine.connection._connections.clear()
+        mongoengine.connection._dbs.clear()
 
     def tearDown(self):
         self.ctx.pop()

tests/test_connection.py

@@ -1,4 +1,7 @@
+import mongoengine
 from mongoengine.context_managers import switch_db
+from nose import SkipTest
+from nose.tools import assert_raises
 import pymongo
 from pymongo.errors import InvalidURI
 from pymongo.read_preferences import ReadPreference
@@ -10,14 +13,21 @@
 
 class ConnectionTestCase(FlaskMongoEngineTestCase):
 
-    def _do_persist(self, db):
+    def _do_persist(self, db, alias=None):
         """Initialize a test Flask application and persist some data in
         MongoDB, ultimately asserting that the connection works.
         """
-        class Todo(db.Document):
-            title = db.StringField(max_length=60)
-            text = db.StringField()
-            done = db.BooleanField(default=False)
+        if alias:
+            class Todo(db.Document):
+                meta = {'db_alias': alias}
+                title = db.StringField(max_length=60)
+                text = db.StringField()
+                done = db.BooleanField(default=False)
+        else:
+            class Todo(db.Document):
+                title = db.StringField(max_length=60)
+                text = db.StringField()
+                done = db.BooleanField(default=False)
 
         db.init_app(self.app)
         Todo.drop_collection()
@@ -41,7 +51,7 @@ def test_simple_connection(self):
             'PORT': 27017,
             'DB': 'flask_mongoengine_test_db'
         }
-        self._do_persist(db)
+        self._do_persist(db, alias='simple_conn')
 
     def test_host_as_uri_string(self):
         """Make sure we can connect to a standalone MongoDB if we specify
@@ -51,14 +61,43 @@ def test_host_as_uri_string(self):
         self.app.config['MONGODB_HOST'] = 'mongodb://localhost:27017/flask_mongoengine_test_db'
         self._do_persist(db)
 
+    def test_mongomock_host_as_uri_string(self):
+        """Make sure we switch to mongomock if we specify the host as a mongomock URI.
+        """
+        if mongoengine.VERSION < (0, 9, 0):
+            raise SkipTest('Mongomock not supported for mongoengine < 0.9.0')
+        db = MongoEngine()
+        self.app.config['MONGODB_HOST'] = 'mongomock://localhost:27017/flask_mongoengine_test_db'
+        with assert_raises(RuntimeError) as exc:
+            self._do_persist(db)
+        assert str(exc.exception) == 'You need mongomock installed to mock MongoEngine.'
+
+    def test_mongomock_as_param(self):
+        """Make sure we switch to mongomock when providing IS_MOCK option.
+        """
+        if mongoengine.VERSION < (0, 9, 0):
+            raise SkipTest('Mongomock not supported for mongoengine < 0.9.0')
+        db = MongoEngine()
+        self.app.config['MONGODB_SETTINGS'] = {
+            'ALIAS': 'simple_conn',
+            'HOST': 'localhost',
+            'PORT': 27017,
+            'DB': 'flask_mongoengine_test_db',
+            'IS_MOCK': True
+        }
+        with assert_raises(RuntimeError) as exc:
+            self._do_persist(db, alias='simple_conn')
+        assert str(exc.exception) == 'You need mongomock installed to mock MongoEngine.'
+
     def test_host_as_list(self):
         """Make sure MONGODB_HOST can be a list hosts."""
         db = MongoEngine()
         self.app.config['MONGODB_SETTINGS'] = {
             'ALIAS': 'host_list',
             'HOST': ['localhost:27017'],
+            'DB': 'flask_mongoengine_test_db'
         }
-        self._do_persist(db)
+        self._do_persist(db, alias='host_list')
 
     def test_multiple_connections(self):
         """Make sure establishing multiple connections to a standalone
@@ -117,6 +156,16 @@ def test_connection_with_invalid_uri(self):
         self.app.config['MONGODB_HOST'] = 'mongo://localhost'
         self.assertRaises(InvalidURI, MongoEngine, self.app)
 
+    def test_ingnored_mongodb_prefix_config(self):
+        """Config starting by MONGODB_ but not used by flask-mongoengine
+        should be ignored.
+        """
+        db = MongoEngine()
+        self.app.config['MONGODB_HOST'] = 'mongodb://localhost:27017/flask_mongoengine_test_db_prod'
+        # Invalid host, should trigger exception if used
+        self.app.config['MONGODB_TEST_HOST'] = 'dummy://localhost:27017/test'
+        self._do_persist(db)
+
     def test_connection_kwargs(self):
         """Make sure additional connection kwargs work."""