b47669403b
This commit: 1. Formats all of our existing code using `black`. 2. Adds a note about using `black` to `CONTRIBUTING.rst`. 3. Runs `black --check` as part of CI (failing builds that aren't properly formatted).
394 lines
12 KiB
Python
394 lines
12 KiB
Python
import datetime
|
|
import re
|
|
import unittest
|
|
|
|
from bson import ObjectId
|
|
|
|
from mongoengine import *
|
|
from mongoengine.errors import InvalidQueryError
|
|
from mongoengine.queryset import Q
|
|
|
|
__all__ = ("QTest",)
|
|
|
|
|
|
class QTest(unittest.TestCase):
|
|
def setUp(self):
|
|
connect(db="mongoenginetest")
|
|
|
|
class Person(Document):
|
|
name = StringField()
|
|
age = IntField()
|
|
meta = {"allow_inheritance": True}
|
|
|
|
Person.drop_collection()
|
|
self.Person = Person
|
|
|
|
def test_empty_q(self):
|
|
"""Ensure that empty Q objects won't hurt.
|
|
"""
|
|
q1 = Q()
|
|
q2 = Q(age__gte=18)
|
|
q3 = Q()
|
|
q4 = Q(name="test")
|
|
q5 = Q()
|
|
|
|
class Person(Document):
|
|
name = StringField()
|
|
age = IntField()
|
|
|
|
query = {"$or": [{"age": {"$gte": 18}}, {"name": "test"}]}
|
|
self.assertEqual((q1 | q2 | q3 | q4 | q5).to_query(Person), query)
|
|
|
|
query = {"age": {"$gte": 18}, "name": "test"}
|
|
self.assertEqual((q1 & q2 & q3 & q4 & q5).to_query(Person), query)
|
|
|
|
def test_q_with_dbref(self):
|
|
"""Ensure Q objects handle DBRefs correctly"""
|
|
connect(db="mongoenginetest")
|
|
|
|
class User(Document):
|
|
pass
|
|
|
|
class Post(Document):
|
|
created_user = ReferenceField(User)
|
|
|
|
user = User.objects.create()
|
|
Post.objects.create(created_user=user)
|
|
|
|
self.assertEqual(Post.objects.filter(created_user=user).count(), 1)
|
|
self.assertEqual(Post.objects.filter(Q(created_user=user)).count(), 1)
|
|
|
|
def test_and_combination(self):
|
|
"""Ensure that Q-objects correctly AND together.
|
|
"""
|
|
|
|
class TestDoc(Document):
|
|
x = IntField()
|
|
y = StringField()
|
|
|
|
query = (Q(x__lt=7) & Q(x__lt=3)).to_query(TestDoc)
|
|
self.assertEqual(query, {"$and": [{"x": {"$lt": 7}}, {"x": {"$lt": 3}}]})
|
|
|
|
query = (Q(y="a") & Q(x__lt=7) & Q(x__lt=3)).to_query(TestDoc)
|
|
self.assertEqual(
|
|
query, {"$and": [{"y": "a"}, {"x": {"$lt": 7}}, {"x": {"$lt": 3}}]}
|
|
)
|
|
|
|
# Check normal cases work without an error
|
|
query = Q(x__lt=7) & Q(x__gt=3)
|
|
|
|
q1 = Q(x__lt=7)
|
|
q2 = Q(x__gt=3)
|
|
query = (q1 & q2).to_query(TestDoc)
|
|
self.assertEqual(query, {"x": {"$lt": 7, "$gt": 3}})
|
|
|
|
# More complex nested example
|
|
query = Q(x__lt=100) & Q(y__ne="NotMyString")
|
|
query &= Q(y__in=["a", "b", "c"]) & Q(x__gt=-100)
|
|
mongo_query = {
|
|
"x": {"$lt": 100, "$gt": -100},
|
|
"y": {"$ne": "NotMyString", "$in": ["a", "b", "c"]},
|
|
}
|
|
self.assertEqual(query.to_query(TestDoc), mongo_query)
|
|
|
|
def test_or_combination(self):
|
|
"""Ensure that Q-objects correctly OR together.
|
|
"""
|
|
|
|
class TestDoc(Document):
|
|
x = IntField()
|
|
|
|
q1 = Q(x__lt=3)
|
|
q2 = Q(x__gt=7)
|
|
query = (q1 | q2).to_query(TestDoc)
|
|
self.assertEqual(query, {"$or": [{"x": {"$lt": 3}}, {"x": {"$gt": 7}}]})
|
|
|
|
def test_and_or_combination(self):
|
|
"""Ensure that Q-objects handle ANDing ORed components.
|
|
"""
|
|
|
|
class TestDoc(Document):
|
|
x = IntField()
|
|
y = BooleanField()
|
|
|
|
TestDoc.drop_collection()
|
|
|
|
query = Q(x__gt=0) | Q(x__exists=False)
|
|
query &= Q(x__lt=100)
|
|
self.assertEqual(
|
|
query.to_query(TestDoc),
|
|
{
|
|
"$and": [
|
|
{"$or": [{"x": {"$gt": 0}}, {"x": {"$exists": False}}]},
|
|
{"x": {"$lt": 100}},
|
|
]
|
|
},
|
|
)
|
|
|
|
q1 = Q(x__gt=0) | Q(x__exists=False)
|
|
q2 = Q(x__lt=100) | Q(y=True)
|
|
query = (q1 & q2).to_query(TestDoc)
|
|
|
|
TestDoc(x=101).save()
|
|
TestDoc(x=10).save()
|
|
TestDoc(y=True).save()
|
|
|
|
self.assertEqual(
|
|
query,
|
|
{
|
|
"$and": [
|
|
{"$or": [{"x": {"$gt": 0}}, {"x": {"$exists": False}}]},
|
|
{"$or": [{"x": {"$lt": 100}}, {"y": True}]},
|
|
]
|
|
},
|
|
)
|
|
self.assertEqual(2, TestDoc.objects(q1 & q2).count())
|
|
|
|
def test_or_and_or_combination(self):
|
|
"""Ensure that Q-objects handle ORing ANDed ORed components. :)
|
|
"""
|
|
|
|
class TestDoc(Document):
|
|
x = IntField()
|
|
y = BooleanField()
|
|
|
|
TestDoc.drop_collection()
|
|
TestDoc(x=-1, y=True).save()
|
|
TestDoc(x=101, y=True).save()
|
|
TestDoc(x=99, y=False).save()
|
|
TestDoc(x=101, y=False).save()
|
|
|
|
q1 = Q(x__gt=0) & (Q(y=True) | Q(y__exists=False))
|
|
q2 = Q(x__lt=100) & (Q(y=False) | Q(y__exists=False))
|
|
query = (q1 | q2).to_query(TestDoc)
|
|
|
|
self.assertEqual(
|
|
query,
|
|
{
|
|
"$or": [
|
|
{
|
|
"$and": [
|
|
{"x": {"$gt": 0}},
|
|
{"$or": [{"y": True}, {"y": {"$exists": False}}]},
|
|
]
|
|
},
|
|
{
|
|
"$and": [
|
|
{"x": {"$lt": 100}},
|
|
{"$or": [{"y": False}, {"y": {"$exists": False}}]},
|
|
]
|
|
},
|
|
]
|
|
},
|
|
)
|
|
self.assertEqual(2, TestDoc.objects(q1 | q2).count())
|
|
|
|
def test_multiple_occurence_in_field(self):
|
|
class Test(Document):
|
|
name = StringField(max_length=40)
|
|
title = StringField(max_length=40)
|
|
|
|
q1 = Q(name__contains="te") | Q(title__contains="te")
|
|
q2 = Q(name__contains="12") | Q(title__contains="12")
|
|
|
|
q3 = q1 & q2
|
|
|
|
query = q3.to_query(Test)
|
|
self.assertEqual(query["$and"][0], q1.to_query(Test))
|
|
self.assertEqual(query["$and"][1], q2.to_query(Test))
|
|
|
|
def test_q_clone(self):
|
|
class TestDoc(Document):
|
|
x = IntField()
|
|
|
|
TestDoc.drop_collection()
|
|
for i in range(1, 101):
|
|
t = TestDoc(x=i)
|
|
t.save()
|
|
|
|
# Check normal cases work without an error
|
|
test = TestDoc.objects(Q(x__lt=7) & Q(x__gt=3))
|
|
|
|
self.assertEqual(test.count(), 3)
|
|
|
|
test2 = test.clone()
|
|
self.assertEqual(test2.count(), 3)
|
|
self.assertNotEqual(test2, test)
|
|
|
|
test3 = test2.filter(x=6)
|
|
self.assertEqual(test3.count(), 1)
|
|
self.assertEqual(test.count(), 3)
|
|
|
|
def test_q(self):
|
|
"""Ensure that Q objects may be used to query for documents.
|
|
"""
|
|
|
|
class BlogPost(Document):
|
|
title = StringField()
|
|
publish_date = DateTimeField()
|
|
published = BooleanField()
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
post1 = BlogPost(
|
|
title="Test 1", publish_date=datetime.datetime(2010, 1, 8), published=False
|
|
)
|
|
post1.save()
|
|
|
|
post2 = BlogPost(
|
|
title="Test 2", publish_date=datetime.datetime(2010, 1, 15), published=True
|
|
)
|
|
post2.save()
|
|
|
|
post3 = BlogPost(title="Test 3", published=True)
|
|
post3.save()
|
|
|
|
post4 = BlogPost(title="Test 4", publish_date=datetime.datetime(2010, 1, 8))
|
|
post4.save()
|
|
|
|
post5 = BlogPost(title="Test 1", publish_date=datetime.datetime(2010, 1, 15))
|
|
post5.save()
|
|
|
|
post6 = BlogPost(title="Test 1", published=False)
|
|
post6.save()
|
|
|
|
# Check ObjectId lookup works
|
|
obj = BlogPost.objects(id=post1.id).first()
|
|
self.assertEqual(obj, post1)
|
|
|
|
# Check Q object combination with one does not exist
|
|
q = BlogPost.objects(Q(title="Test 5") | Q(published=True))
|
|
posts = [post.id for post in q]
|
|
|
|
published_posts = (post2, post3)
|
|
self.assertTrue(all(obj.id in posts for obj in published_posts))
|
|
|
|
q = BlogPost.objects(Q(title="Test 1") | Q(published=True))
|
|
posts = [post.id for post in q]
|
|
published_posts = (post1, post2, post3, post5, post6)
|
|
self.assertTrue(all(obj.id in posts for obj in published_posts))
|
|
|
|
# Check Q object combination
|
|
date = datetime.datetime(2010, 1, 10)
|
|
q = BlogPost.objects(Q(publish_date__lte=date) | Q(published=True))
|
|
posts = [post.id for post in q]
|
|
|
|
published_posts = (post1, post2, post3, post4)
|
|
self.assertTrue(all(obj.id in posts for obj in published_posts))
|
|
|
|
self.assertFalse(any(obj.id in posts for obj in [post5, post6]))
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
# Check the 'in' operator
|
|
self.Person(name="user1", age=20).save()
|
|
self.Person(name="user2", age=20).save()
|
|
self.Person(name="user3", age=30).save()
|
|
self.Person(name="user4", age=40).save()
|
|
|
|
self.assertEqual(self.Person.objects(Q(age__in=[20])).count(), 2)
|
|
self.assertEqual(self.Person.objects(Q(age__in=[20, 30])).count(), 3)
|
|
|
|
# Test invalid query objs
|
|
with self.assertRaises(InvalidQueryError):
|
|
self.Person.objects("user1")
|
|
|
|
# filter should fail, too
|
|
with self.assertRaises(InvalidQueryError):
|
|
self.Person.objects.filter("user1")
|
|
|
|
def test_q_regex(self):
|
|
"""Ensure that Q objects can be queried using regexes.
|
|
"""
|
|
person = self.Person(name="Guido van Rossum")
|
|
person.save()
|
|
|
|
obj = self.Person.objects(Q(name=re.compile("^Gui"))).first()
|
|
self.assertEqual(obj, person)
|
|
obj = self.Person.objects(Q(name=re.compile("^gui"))).first()
|
|
self.assertEqual(obj, None)
|
|
|
|
obj = self.Person.objects(Q(name=re.compile("^gui", re.I))).first()
|
|
self.assertEqual(obj, person)
|
|
|
|
obj = self.Person.objects(Q(name__not=re.compile("^bob"))).first()
|
|
self.assertEqual(obj, person)
|
|
|
|
obj = self.Person.objects(Q(name__not=re.compile("^Gui"))).first()
|
|
self.assertEqual(obj, None)
|
|
|
|
def test_q_repr(self):
|
|
self.assertEqual(repr(Q()), "Q(**{})")
|
|
self.assertEqual(repr(Q(name="test")), "Q(**{'name': 'test'})")
|
|
|
|
self.assertEqual(
|
|
repr(Q(name="test") & Q(age__gte=18)),
|
|
"(Q(**{'name': 'test'}) & Q(**{'age__gte': 18}))",
|
|
)
|
|
|
|
self.assertEqual(
|
|
repr(Q(name="test") | Q(age__gte=18)),
|
|
"(Q(**{'name': 'test'}) | Q(**{'age__gte': 18}))",
|
|
)
|
|
|
|
def test_q_lists(self):
|
|
"""Ensure that Q objects query ListFields correctly.
|
|
"""
|
|
|
|
class BlogPost(Document):
|
|
tags = ListField(StringField())
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
BlogPost(tags=["python", "mongo"]).save()
|
|
BlogPost(tags=["python"]).save()
|
|
|
|
self.assertEqual(BlogPost.objects(Q(tags="mongo")).count(), 1)
|
|
self.assertEqual(BlogPost.objects(Q(tags="python")).count(), 2)
|
|
|
|
BlogPost.drop_collection()
|
|
|
|
def test_q_merge_queries_edge_case(self):
|
|
class User(Document):
|
|
email = EmailField(required=False)
|
|
name = StringField()
|
|
|
|
User.drop_collection()
|
|
pk = ObjectId()
|
|
User(email="example@example.com", pk=pk).save()
|
|
|
|
self.assertEqual(
|
|
1,
|
|
User.objects.filter(Q(email="example@example.com") | Q(name="John Doe"))
|
|
.limit(2)
|
|
.filter(pk=pk)
|
|
.count(),
|
|
)
|
|
|
|
def test_chained_q_or_filtering(self):
|
|
class Post(EmbeddedDocument):
|
|
name = StringField(required=True)
|
|
|
|
class Item(Document):
|
|
postables = ListField(EmbeddedDocumentField(Post))
|
|
|
|
Item.drop_collection()
|
|
|
|
Item(postables=[Post(name="a"), Post(name="b")]).save()
|
|
Item(postables=[Post(name="a"), Post(name="c")]).save()
|
|
Item(postables=[Post(name="a"), Post(name="b"), Post(name="c")]).save()
|
|
|
|
self.assertEqual(
|
|
Item.objects(Q(postables__name="a") & Q(postables__name="b")).count(), 2
|
|
)
|
|
self.assertEqual(
|
|
Item.objects.filter(postables__name="a")
|
|
.filter(postables__name="b")
|
|
.count(),
|
|
2,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|