Merge pull request MongoEngine#2732 from bagerard/davidlatwe-alternative_conection

[Clone] Allow specifying alternative connection_class

Author: bagerard

docs/changelog.rst

@@ -7,6 +7,10 @@ Changelog
 Development
 ===========
 - (Fill this out as you fix issues and develop your features).
+- Added `mongo_client_class` optional parameter to connect() to allow to use an alternative mongo client than pymongo.MongoClient.
+  Typically to support mock mongo libraries like mongomock, montydb, mongita #2729
+- BREAKING CHANGE: connecting MongoEngine with mongomock should now use the new `mongo_client_class`
+  For more info, check https://docs.mongoengine.org/guide/mongomock.html
 
 Changes in 0.26.0
 =================

docs/guide/mongomock.rst

@@ -1,23 +1,28 @@
-==============================
+=========================
 Use mongomock for testing
-==============================
+=========================
 
-`mongomock <https://github.com/vmalloc/mongomock/>`_ is a package to do just
-what the name implies, mocking a mongo database.
+Although we recommend running your tests against a regular MongoDB server, it is sometimes useful to plug
+MongoEngine to alternative implementations (mongomock, montydb, mongita, etc).
+
+`mongomock <https://github.com/vmalloc/mongomock/>`_ is historically the one suggested for MongoEngine and is
+a package to do just what the name implies, mocking a mongo database.
 
 To use with mongoengine, simply specify mongomock when connecting with
 mongoengine:
 
 .. code-block:: python
 
-    connect('mongoenginetest', host='mongomock://localhost')
+    import mongomock
+
+    connect('mongoenginetest', host='mongodb://localhost', mongo_client_class=mongomock.MongoClient)
     conn = get_connection()
 
 or with an alias:
 
 .. code-block:: python
 
-    connect('mongoenginetest', host='mongomock://localhost', alias='testdb')
+    connect('mongoenginetest', host='mongodb://localhost', mongo_client_class=mongomock.MongoClient, alias='testdb')
     conn = get_connection('testdb')
 
 Example of test file:
@@ -34,7 +39,7 @@ Example of test file:
 
         @classmethod
         def setUpClass(cls):
-            connect('mongoenginetest', host='mongomock://localhost')
+            connect('mongoenginetest', host='mongodb://localhost', mongo_client_class=mongomock.MongoClient)
 
         @classmethod
         def tearDownClass(cls):

mongoengine/connection.py

@@ -74,8 +74,9 @@ def _get_connection_settings(
     :param authentication_mechanism: database authentication mechanisms.
         By default, use SCRAM-SHA-1 with MongoDB 3.0 and later,
         MONGODB-CR (MongoDB Challenge Response protocol) for older servers.
-    :param is_mock: explicitly use mongomock for this connection
-        (can also be done by using `mongomock: // ` as db host prefix)
+    :param mongo_client_class: using alternative connection client other than
+        pymongo.MongoClient, e.g. mongomock, montydb, that provides pymongo alike
+        interface but not necessarily for connecting to a real mongo instance.
     :param kwargs: ad-hoc parameters to be passed into the pymongo driver,
         for example maxpoolsize, tz_aware, etc. See the documentation
         for pymongo's `MongoClient` for a full list.
@@ -102,22 +103,17 @@ def _get_connection_settings(
     resolved_hosts = []
     for entity in conn_host:
 
-        # Handle Mongomock
-        if entity.startswith("mongomock://"):
-            conn_settings["is_mock"] = True
-            # `mongomock://` is not a valid url prefix and must be replaced by `mongodb://`
-            new_entity = entity.replace("mongomock://", "mongodb://", 1)
-            resolved_hosts.append(new_entity)
-
-            uri_dict = uri_parser.parse_uri(new_entity)
-
-            database = uri_dict.get("database")
-            if database:
-                conn_settings["name"] = database
+        # Reject old mongomock integration
+        # To be removed in a few versions after 0.27.0
+        if entity.startswith("mongomock://") or kwargs.get("is_mock"):
+            raise Exception(
+                "Use of mongomock:// URI or 'is_mock' were removed in favor of 'mongo_client_class=mongomock.MongoClient'. "
+                "Check the CHANGELOG for more info"
+            )
 
         # Handle URI style connections, only updating connection params which
         # were explicitly specified in the URI.
-        elif "://" in entity:
+        if "://" in entity:
             uri_dict = uri_parser.parse_uri(entity)
             resolved_hosts.append(entity)
 
@@ -219,8 +215,9 @@ def register_connection(
     :param authentication_mechanism: database authentication mechanisms.
         By default, use SCRAM-SHA-1 with MongoDB 3.0 and later,
         MONGODB-CR (MongoDB Challenge Response protocol) for older servers.
-    :param is_mock: explicitly use mongomock for this connection
-        (can also be done by using `mongomock: // ` as db host prefix)
+    :param mongo_client_class: using alternative connection client other than
+        pymongo.MongoClient, e.g. mongomock, montydb, that provides pymongo alike
+        interface but not necessarily for connecting to a real mongo instance.
     :param kwargs: ad-hoc parameters to be passed into the pymongo driver,
         for example maxpoolsize, tz_aware, etc. See the documentation
         for pymongo's `MongoClient` for a full list.
@@ -326,35 +323,30 @@ def _clean_settings(settings_dict):
     conn_settings = _clean_settings(raw_conn_settings)
 
     # Determine if we should use PyMongo's or mongomock's MongoClient.
-    is_mock = conn_settings.pop("is_mock", False)
-    if is_mock:
-        try:
-            import mongomock
-        except ImportError:
-            raise RuntimeError("You need mongomock installed to mock MongoEngine.")
-        connection_class = mongomock.MongoClient
+    if "mongo_client_class" in conn_settings:
+        mongo_client_class = conn_settings.pop("mongo_client_class")
     else:
-        connection_class = MongoClient
+        mongo_client_class = MongoClient
 
     # Re-use existing connection if one is suitable.
     existing_connection = _find_existing_connection(raw_conn_settings)
     if existing_connection:
         connection = existing_connection
     else:
         connection = _create_connection(
-            alias=alias, connection_class=connection_class, **conn_settings
+            alias=alias, mongo_client_class=mongo_client_class, **conn_settings
         )
     _connections[alias] = connection
     return _connections[alias]
 
 
-def _create_connection(alias, connection_class, **connection_settings):
+def _create_connection(alias, mongo_client_class, **connection_settings):
     """
     Create the new connection for this alias. Raise
     ConnectionFailure if it can't be established.
     """
     try:
-        return connection_class(**connection_settings)
+        return mongo_client_class(**connection_settings)
     except Exception as e:
         raise ConnectionFailure(f"Cannot connect to database {alias} :\n{e}")
 

tests/test_connection.py

@@ -53,15 +53,21 @@ def test_connect(self):
         connect("mongoenginetest")
 
         conn = get_connection()
-        assert isinstance(conn, pymongo.mongo_client.MongoClient)
+        assert isinstance(conn, pymongo.MongoClient)
 
         db = get_db()
         assert isinstance(db, pymongo.database.Database)
         assert db.name == "mongoenginetest"
 
         connect("mongoenginetest2", alias="testdb")
         conn = get_connection("testdb")
-        assert isinstance(conn, pymongo.mongo_client.MongoClient)
+        assert isinstance(conn, pymongo.MongoClient)
+
+        connect(
+            "mongoenginetest2", alias="testdb3", mongo_client_class=pymongo.MongoClient
+        )
+        conn = get_connection("testdb")
+        assert isinstance(conn, pymongo.MongoClient)
 
     def test_connect_disconnect_works_properly(self):
         class History1(Document):

tests/test_connection_mongomock.py

@@ -32,45 +32,76 @@ def tearDown(self):
         mongoengine.connection._connections = {}
         mongoengine.connection._dbs = {}
 
+    @require_mongomock
+    def test_connect_raise_if_mongomock_uri_provided(self):
+        with pytest.raises(
+            Exception, match="Use of mongomock:// URI or 'is_mock' were removed"
+        ):
+            connect("test", host="mongomock://localhost")
+
+    @require_mongomock
+    def test_connect_raise_if_is_mock_provided(self):
+        with pytest.raises(
+            Exception, match="Use of mongomock:// URI or 'is_mock' were removed"
+        ):
+            connect("test", host="mongodb://localhost", is_mock=True)
+
     @require_mongomock
     def test_connect_in_mocking(self):
         """Ensure that the connect() method works properly in mocking."""
-        connect("mongoenginetest", host="mongomock://localhost")
+        connect(
+            "mongoenginetest",
+            host="mongodb://localhost",
+            mongo_client_class=mongomock.MongoClient,
+        )
         conn = get_connection()
         assert isinstance(conn, mongomock.MongoClient)
 
-        connect("mongoenginetest2", host="mongomock://localhost", alias="testdb2")
+        connect(
+            "mongoenginetest2",
+            host="mongodb://localhost",
+            mongo_client_class=mongomock.MongoClient,
+            alias="testdb2",
+        )
         conn = get_connection("testdb2")
         assert isinstance(conn, mongomock.MongoClient)
 
         connect(
             "mongoenginetest3",
             host="mongodb://localhost",
-            is_mock=True,
+            mongo_client_class=mongomock.MongoClient,
             alias="testdb3",
         )
         conn = get_connection("testdb3")
         assert isinstance(conn, mongomock.MongoClient)
 
-        connect("mongoenginetest4", is_mock=True, alias="testdb4")
+        connect(
+            "mongoenginetest4",
+            mongo_client_class=mongomock.MongoClient,
+            alias="testdb4",
+        )
         conn = get_connection("testdb4")
         assert isinstance(conn, mongomock.MongoClient)
 
         connect(
             host="mongodb://localhost:27017/mongoenginetest5",
-            is_mock=True,
+            mongo_client_class=mongomock.MongoClient,
             alias="testdb5",
         )
         conn = get_connection("testdb5")
         assert isinstance(conn, mongomock.MongoClient)
 
-        connect(host="mongomock://localhost:27017/mongoenginetest6", alias="testdb6")
+        connect(
+            host="mongodb://localhost:27017/mongoenginetest6",
+            mongo_client_class=mongomock.MongoClient,
+            alias="testdb6",
+        )
         conn = get_connection("testdb6")
         assert isinstance(conn, mongomock.MongoClient)
 
         connect(
-            host="mongomock://localhost:27017/mongoenginetest7",
-            is_mock=True,
+            host="mongodb://localhost:27017/mongoenginetest7",
+            mongo_client_class=mongomock.MongoClient,
             alias="testdb7",
         )
         conn = get_connection("testdb7")
@@ -84,7 +115,10 @@ def test_default_database_with_mocking(self):
         class SomeDocument(Document):
             pass
 
-        conn = connect(host="mongomock://localhost:27017/mongoenginetest")
+        conn = connect(
+            host="mongodb://localhost:27017/mongoenginetest",
+            mongo_client_class=mongomock.MongoClient,
+        )
         some_document = SomeDocument()
         # database won't exist until we save a document
         some_document.save()
@@ -96,7 +130,10 @@ class SomeDocument(Document):
     def test_basic_queries_against_mongomock(self):
         disconnect_all()
 
-        connect(host="mongomock://localhost:27017/mongoenginetest")
+        connect(
+            host="mongodb://localhost:27017/mongoenginetest",
+            mongo_client_class=mongomock.MongoClient,
+        )
 
         class Person(Document):
             name = StringField()
@@ -129,35 +166,38 @@ def test_connect_with_host_list(self):
 
         Uses mongomock to test w/o needing multiple mongod/mongos processes
         """
-        connect(host=["mongomock://localhost"])
+        connect(host=["mongodb://localhost"], mongo_client_class=mongomock.MongoClient)
         conn = get_connection()
         assert isinstance(conn, mongomock.MongoClient)
 
-        connect(host=["mongodb://localhost"], is_mock=True, alias="testdb2")
-        conn = get_connection("testdb2")
-        assert isinstance(conn, mongomock.MongoClient)
-
-        connect(host=["localhost"], is_mock=True, alias="testdb3")
+        connect(
+            host=["localhost"],
+            mongo_client_class=mongomock.MongoClient,
+            alias="testdb3",
+        )
         conn = get_connection("testdb3")
         assert isinstance(conn, mongomock.MongoClient)
 
         connect(
-            host=["mongomock://localhost:27017", "mongomock://localhost:27018"],
+            host=["mongodb://localhost:27017", "mongodb://localhost:27018"],
             alias="testdb4",
+            mongo_client_class=mongomock.MongoClient,
         )
         conn = get_connection("testdb4")
         assert isinstance(conn, mongomock.MongoClient)
 
         connect(
             host=["mongodb://localhost:27017", "mongodb://localhost:27018"],
-            is_mock=True,
+            mongo_client_class=mongomock.MongoClient,
             alias="testdb5",
         )
         conn = get_connection("testdb5")
         assert isinstance(conn, mongomock.MongoClient)
 
         connect(
-            host=["localhost:27017", "localhost:27018"], is_mock=True, alias="testdb6"
+            host=["localhost:27017", "localhost:27018"],
+            mongo_client_class=mongomock.MongoClient,
+            alias="testdb6",
         )
         conn = get_connection("testdb6")
         assert isinstance(conn, mongomock.MongoClient)