Skip to content

Commit ccaa42d

Browse files
bagerardbagerard
committed
add tests for recent propagation of comment/hint/collation to delete/update/aggregate
1 parent 199ddd1 commit ccaa42d

File tree

5 files changed

+173
-13
lines changed

5 files changed

+173
-13
lines changed

‎tests/document/test_instance.py‎

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,11 @@
4242
PickleSignalsTest,
4343
PickleTest,
4444
)
45-
from tests.utils import MongoDBTestCase, get_as_pymongo
45+
from tests.utils import (
46+
MongoDBTestCase,
47+
db_ops_tracker,
48+
get_as_pymongo,
49+
)
4650

4751
TEST_IMAGE_PATH = os.path.join(os.path.dirname(__file__), "../fields/mongoengine.png")
4852

@@ -1879,6 +1883,49 @@ def test_delete(self):
18791883
person.delete()
18801884
assert self.Person.objects.count() == 0
18811885

1886+
def test_delete_propagates_hint_collation_and_comment(self):
1887+
"""Make sure adding a hint/comment/collation to the query gets added to the query"""
1888+
mongo_ver = get_mongodb_version()
1889+
1890+
base = {"locale": "en", "strength": 2}
1891+
index_name = "name_1"
1892+
1893+
class AggPerson(Document):
1894+
name = StringField()
1895+
meta = {
1896+
"indexes": [{"fields": ["name"], "name": index_name, "collation": base}]
1897+
}
1898+
1899+
AggPerson.drop_collection()
1900+
_ = AggPerson.objects.first()
1901+
1902+
comment = "test_comment"
1903+
1904+
with db_ops_tracker() as q:
1905+
_ = AggPerson.objects().comment(comment).delete()
1906+
query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
1907+
CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
1908+
assert "hint" not in query_op[CMD_QUERY_KEY]
1909+
assert query_op[CMD_QUERY_KEY]["comment"] == comment
1910+
assert "collation" not in query_op[CMD_QUERY_KEY]
1911+
1912+
with db_ops_tracker() as q:
1913+
_ = AggPerson.objects.hint(index_name).delete()
1914+
query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
1915+
CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
1916+
1917+
assert query_op[CMD_QUERY_KEY]["hint"] == {"$hint": index_name}
1918+
assert "comment" not in query_op[CMD_QUERY_KEY]
1919+
assert "collation" not in query_op[CMD_QUERY_KEY]
1920+
1921+
with db_ops_tracker() as q:
1922+
_ = AggPerson.objects.collation(base).delete()
1923+
query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
1924+
CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
1925+
assert "hint" not in query_op[CMD_QUERY_KEY]
1926+
assert "comment" not in query_op[CMD_QUERY_KEY]
1927+
assert query_op[CMD_QUERY_KEY]["collation"] == base
1928+
18821929
def test_save_custom_id(self):
18831930
"""Ensure that a document may be saved with a custom _id."""
18841931

‎tests/fields/test_binary_field.py‎

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,11 @@
44
from bson import Binary
55

66
from mongoengine import *
7-
from tests.utils import MongoDBTestCase
7+
from mongoengine.mongodb_support import (
8+
MONGODB_36,
9+
get_mongodb_version,
10+
)
11+
from tests.utils import MongoDBTestCase, db_ops_tracker
812

913
BIN_VALUE = "\xa9\xf3\x8d(\xd7\x03\x84\xb4k[\x0f\xe3\xa2\x19\x85p[J\xa3\xd2>\xde\xe6\x87\xb1\x7f\xc6\xe6\xd9r\x18\xf5".encode(
1014
"latin-1"
@@ -146,3 +150,46 @@ class MyDocument(Document):
146150
assert n_updated == 1
147151
fetched = MyDocument.objects.with_id(doc.id)
148152
assert fetched.bin_field == BIN_VALUE
153+
154+
def test_update_propagates_hint_collation_and_comment(self):
155+
"""Make sure adding a hint/comment/collation to the query gets added to the query"""
156+
mongo_ver = get_mongodb_version()
157+
158+
base = {"locale": "en", "strength": 2}
159+
index_name = "name_1"
160+
161+
class AggPerson(Document):
162+
name = StringField()
163+
meta = {
164+
"indexes": [{"fields": ["name"], "name": index_name, "collation": base}]
165+
}
166+
167+
AggPerson.drop_collection()
168+
_ = AggPerson.objects.first()
169+
170+
comment = "test_comment"
171+
172+
with db_ops_tracker() as q:
173+
_ = AggPerson.objects.comment(comment).update_one(name="something")
174+
query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
175+
CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
176+
assert "hint" not in query_op[CMD_QUERY_KEY]
177+
assert query_op[CMD_QUERY_KEY]["comment"] == comment
178+
assert "collation" not in query_op[CMD_QUERY_KEY]
179+
180+
with db_ops_tracker() as q:
181+
_ = AggPerson.objects.hint(index_name).update_one(name="something")
182+
query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
183+
CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
184+
185+
assert query_op[CMD_QUERY_KEY]["hint"] == {"$hint": index_name}
186+
assert "comment" not in query_op[CMD_QUERY_KEY]
187+
assert "collation" not in query_op[CMD_QUERY_KEY]
188+
189+
with db_ops_tracker() as q:
190+
_ = AggPerson.objects.collation(base).update_one(name="something")
191+
query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
192+
CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
193+
assert "hint" not in query_op[CMD_QUERY_KEY]
194+
assert "comment" not in query_op[CMD_QUERY_KEY]
195+
assert query_op[CMD_QUERY_KEY]["collation"] == base

‎tests/queryset/test_queryset.py‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -27,21 +27,13 @@
2727
)
2828
from mongoengine.queryset.base import BaseQuerySet
2929
from tests.utils import (
30+
db_ops_tracker,
3031
requires_mongodb_gte_42,
3132
requires_mongodb_gte_44,
3233
requires_mongodb_lt_42,
3334
)
3435

3536

36-
class db_ops_tracker(query_counter):
37-
def get_ops(self):
38-
ignore_query = dict(self._ignored_query)
39-
ignore_query["command.count"] = {
40-
"$ne": "system.profile"
41-
} # Ignore the query issued by query_counter
42-
return list(self.db.system.profile.find(ignore_query))
43-
44-
4537
def get_key_compat(mongo_ver):
4638
ORDER_BY_KEY = "sort"
4739
CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"

‎tests/queryset/test_queryset_aggregation.py‎

Lines changed: 66 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,12 @@
33

44
from pymongo.read_preferences import ReadPreference
55

6-
from mongoengine import *
7-
from tests.utils import MongoDBTestCase
6+
from mongoengine import Document, IntField, PointField, StringField
7+
from mongoengine.mongodb_support import (
8+
MONGODB_36,
9+
get_mongodb_version,
10+
)
11+
from tests.utils import MongoDBTestCase, db_ops_tracker
812

913

1014
class TestQuerysetAggregate(MongoDBTestCase):
@@ -87,6 +91,66 @@ class Person(Document):
8791
{"_id": p3.pk, "name": "SANDRA MARA"},
8892
]
8993

94+
def test_aggregation_comment(self):
95+
"""Make sure adding a comment to the query gets added to the query"""
96+
mongo_ver = get_mongodb_version()
97+
98+
class AggPerson(Document):
99+
name = StringField()
100+
101+
AggPerson.drop_collection()
102+
103+
pipeline = [{"$project": {"name": {"$toUpper": "$name"}}}]
104+
comment = "some_comment"
105+
with db_ops_tracker() as q:
106+
_ = list(AggPerson.objects.comment(comment).aggregate(pipeline))
107+
query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
108+
CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
109+
assert query_op[CMD_QUERY_KEY]["comment"] == comment
110+
111+
def test_aggregation_propagates_hint_collation_and_comment(self):
112+
"""Make sure adding a hint/comment/collation to the query gets added to the query"""
113+
mongo_ver = get_mongodb_version()
114+
115+
base = {"locale": "en", "strength": 2}
116+
index_name = "name_1"
117+
118+
class AggPerson(Document):
119+
name = StringField()
120+
meta = {
121+
"indexes": [{"fields": ["name"], "name": index_name, "collation": base}]
122+
}
123+
124+
AggPerson.drop_collection()
125+
_ = AggPerson.objects.first()
126+
127+
pipeline = [{"$project": {"name": {"$toUpper": "$name"}}}]
128+
comment = "test_comment"
129+
130+
with db_ops_tracker() as q:
131+
_ = list(AggPerson.objects.comment(comment).aggregate(pipeline))
132+
query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
133+
CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
134+
assert "hint" not in query_op[CMD_QUERY_KEY]
135+
assert query_op[CMD_QUERY_KEY]["comment"] == comment
136+
assert "collation" not in query_op[CMD_QUERY_KEY]
137+
138+
with db_ops_tracker() as q:
139+
_ = list(AggPerson.objects.hint(index_name).aggregate(pipeline))
140+
query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
141+
CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
142+
assert query_op[CMD_QUERY_KEY]["hint"] == "name_1"
143+
assert "comment" not in query_op[CMD_QUERY_KEY]
144+
assert "collation" not in query_op[CMD_QUERY_KEY]
145+
146+
with db_ops_tracker() as q:
147+
_ = list(AggPerson.objects.collation(base).aggregate(pipeline))
148+
query_op = q.db.system.profile.find({"ns": "mongoenginetest.agg_person"})[0]
149+
CMD_QUERY_KEY = "command" if mongo_ver >= MONGODB_36 else "query"
150+
assert "hint" not in query_op[CMD_QUERY_KEY]
151+
assert "comment" not in query_op[CMD_QUERY_KEY]
152+
assert query_op[CMD_QUERY_KEY]["collation"] == base
153+
90154
def test_queryset_aggregation_with_limit(self):
91155
class Person(Document):
92156
name = StringField()

‎tests/utils.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
from mongoengine import connect
88
from mongoengine.connection import disconnect_all, get_db
9+
from mongoengine.context_managers import query_counter
910
from mongoengine.mongodb_support import get_mongodb_version
1011

1112
MONGO_TEST_DB = "mongoenginetest" # standard name for the test database
@@ -92,3 +93,12 @@ def _inner(*args, **kwargs):
9293
pytest.skip(f"Needs MongoDB {oper.__name__} v{pretty_version}")
9394

9495
return _inner
96+
97+
98+
class db_ops_tracker(query_counter):
99+
def get_ops(self):
100+
ignore_query = dict(self._ignored_query)
101+
ignore_query["command.count"] = {
102+
"$ne": "system.profile"
103+
} # Ignore the query issued by query_counter
104+
return list(self.db.system.profile.find(ignore_query))

0 commit comments

Comments
 (0)