|
| 1 | +import unittest |
| 2 | + |
| 3 | +import pytest |
| 4 | + |
| 5 | +import mongoengine.connection |
| 6 | +from mongoengine import ( |
| 7 | + Document, |
| 8 | + StringField, |
| 9 | + connect, |
| 10 | + disconnect_all, |
| 11 | +) |
| 12 | +from mongoengine.connection import get_connection |
| 13 | + |
| 14 | + |
| 15 | +try: |
| 16 | + import mongomock |
| 17 | + |
| 18 | + MONGOMOCK_INSTALLED = True |
| 19 | +except ImportError: |
| 20 | + MONGOMOCK_INSTALLED = False |
| 21 | + |
| 22 | +require_mongomock = pytest.mark.skipif( |
| 23 | + not MONGOMOCK_INSTALLED, reason="you need mongomock installed to run this testcase" |
| 24 | +) |
| 25 | + |
| 26 | + |
| 27 | +class MongoMockConnectionTest(unittest.TestCase): |
| 28 | + @classmethod |
| 29 | + def setUpClass(cls): |
| 30 | + disconnect_all() |
| 31 | + |
| 32 | + @classmethod |
| 33 | + def tearDownClass(cls): |
| 34 | + disconnect_all() |
| 35 | + |
| 36 | + def tearDown(self): |
| 37 | + mongoengine.connection._connection_settings = {} |
| 38 | + mongoengine.connection._connections = {} |
| 39 | + mongoengine.connection._dbs = {} |
| 40 | + |
| 41 | + @require_mongomock |
| 42 | + def test_connect_in_mocking(self): |
| 43 | + """Ensure that the connect() method works properly in mocking.""" |
| 44 | + connect("mongoenginetest", host="mongomock://localhost") |
| 45 | + conn = get_connection() |
| 46 | + assert isinstance(conn, mongomock.MongoClient) |
| 47 | + |
| 48 | + connect("mongoenginetest2", host="mongomock://localhost", alias="testdb2") |
| 49 | + conn = get_connection("testdb2") |
| 50 | + assert isinstance(conn, mongomock.MongoClient) |
| 51 | + |
| 52 | + connect( |
| 53 | + "mongoenginetest3", |
| 54 | + host="mongodb://localhost", |
| 55 | + is_mock=True, |
| 56 | + alias="testdb3", |
| 57 | + ) |
| 58 | + conn = get_connection("testdb3") |
| 59 | + assert isinstance(conn, mongomock.MongoClient) |
| 60 | + |
| 61 | + connect("mongoenginetest4", is_mock=True, alias="testdb4") |
| 62 | + conn = get_connection("testdb4") |
| 63 | + assert isinstance(conn, mongomock.MongoClient) |
| 64 | + |
| 65 | + connect( |
| 66 | + host="mongodb://localhost:27017/mongoenginetest5", |
| 67 | + is_mock=True, |
| 68 | + alias="testdb5", |
| 69 | + ) |
| 70 | + conn = get_connection("testdb5") |
| 71 | + assert isinstance(conn, mongomock.MongoClient) |
| 72 | + |
| 73 | + connect(host="mongomock://localhost:27017/mongoenginetest6", alias="testdb6") |
| 74 | + conn = get_connection("testdb6") |
| 75 | + assert isinstance(conn, mongomock.MongoClient) |
| 76 | + |
| 77 | + connect( |
| 78 | + host="mongomock://localhost:27017/mongoenginetest7", |
| 79 | + is_mock=True, |
| 80 | + alias="testdb7", |
| 81 | + ) |
| 82 | + conn = get_connection("testdb7") |
| 83 | + assert isinstance(conn, mongomock.MongoClient) |
| 84 | + |
| 85 | + @require_mongomock |
| 86 | + def test_default_database_with_mocking(self): |
| 87 | + """Ensure that the default database is correctly set when using mongomock.""" |
| 88 | + disconnect_all() |
| 89 | + |
| 90 | + class SomeDocument(Document): |
| 91 | + pass |
| 92 | + |
| 93 | + conn = connect(host="mongomock://localhost:27017/mongoenginetest") |
| 94 | + some_document = SomeDocument() |
| 95 | + # database won't exist until we save a document |
| 96 | + some_document.save() |
| 97 | + assert SomeDocument.objects.count() == 1 |
| 98 | + assert conn.get_default_database().name == "mongoenginetest" |
| 99 | + assert conn.list_database_names()[0] == "mongoenginetest" |
| 100 | + |
| 101 | + @require_mongomock |
| 102 | + def test_basic_queries_against_mongomock(self): |
| 103 | + disconnect_all() |
| 104 | + |
| 105 | + connect(host="mongomock://localhost:27017/mongoenginetest") |
| 106 | + |
| 107 | + class Person(Document): |
| 108 | + name = StringField() |
| 109 | + |
| 110 | + Person.drop_collection() |
| 111 | + assert Person.objects.count() == 0 |
| 112 | + |
| 113 | + bob = Person(name="Bob").save() |
| 114 | + john = Person(name="John").save() |
| 115 | + assert Person.objects.count() == 2 |
| 116 | + |
| 117 | + qs = Person.objects(name="Bob") |
| 118 | + assert qs.count() == 1 |
| 119 | + assert qs.first() == bob |
| 120 | + assert list(qs.as_pymongo()) == [{"_id": bob.id, "name": "Bob"}] |
| 121 | + |
| 122 | + pipeline = [{"$project": {"name": {"$toUpper": "$name"}}}] |
| 123 | + data = Person.objects.order_by("name").aggregate(pipeline) |
| 124 | + assert list(data) == [ |
| 125 | + {"_id": bob.id, "name": "BOB"}, |
| 126 | + {"_id": john.id, "name": "JOHN"}, |
| 127 | + ] |
| 128 | + |
| 129 | + Person.drop_collection() |
| 130 | + assert Person.objects.count() == 0 |
| 131 | + |
| 132 | + @require_mongomock |
| 133 | + def test_connect_with_host_list(self): |
| 134 | + """Ensure that the connect() method works when host is a list |
| 135 | +
|
| 136 | + Uses mongomock to test w/o needing multiple mongod/mongos processes |
| 137 | + """ |
| 138 | + connect(host=["mongomock://localhost"]) |
| 139 | + conn = get_connection() |
| 140 | + assert isinstance(conn, mongomock.MongoClient) |
| 141 | + |
| 142 | + connect(host=["mongodb://localhost"], is_mock=True, alias="testdb2") |
| 143 | + conn = get_connection("testdb2") |
| 144 | + assert isinstance(conn, mongomock.MongoClient) |
| 145 | + |
| 146 | + connect(host=["localhost"], is_mock=True, alias="testdb3") |
| 147 | + conn = get_connection("testdb3") |
| 148 | + assert isinstance(conn, mongomock.MongoClient) |
| 149 | + |
| 150 | + connect( |
| 151 | + host=["mongomock://localhost:27017", "mongomock://localhost:27018"], |
| 152 | + alias="testdb4", |
| 153 | + ) |
| 154 | + conn = get_connection("testdb4") |
| 155 | + assert isinstance(conn, mongomock.MongoClient) |
| 156 | + |
| 157 | + connect( |
| 158 | + host=["mongodb://localhost:27017", "mongodb://localhost:27018"], |
| 159 | + is_mock=True, |
| 160 | + alias="testdb5", |
| 161 | + ) |
| 162 | + conn = get_connection("testdb5") |
| 163 | + assert isinstance(conn, mongomock.MongoClient) |
| 164 | + |
| 165 | + connect( |
| 166 | + host=["localhost:27017", "localhost:27018"], is_mock=True, alias="testdb6" |
| 167 | + ) |
| 168 | + conn = get_connection("testdb6") |
| 169 | + assert isinstance(conn, mongomock.MongoClient) |
| 170 | + |
| 171 | + |
| 172 | +if __name__ == "__main__": |
| 173 | + unittest.main() |
0 commit comments