diff --git a/tests/lookup_/tests.py b/tests/lookup_/tests.py index 6166e1003..b6ac8a322 100644 --- a/tests/lookup_/tests.py +++ b/tests/lookup_/tests.py @@ -1,4 +1,5 @@ from bson import SON +from django.db.models import Sum from django.test import TestCase from django_mongodb_backend.test import MongoTestCaseMixin @@ -73,3 +74,112 @@ def test_eq_and_in(self): "lookup__book", [{"$match": {"$and": [{"isbn": {"$in": ("12345", "56789")}}, {"title": "Moby Dick"}]}}], ) + + def test_gt(self): + with self.assertNumQueries(1) as ctx: + list(Number.objects.filter(num__gt=2)) + self.assertAggregateQuery( + ctx.captured_queries[0]["sql"], + "lookup__number", + [ + {"$match": {"num": {"$gt": 2}}}, + {"$addFields": {"num": "$num"}}, + {"$sort": SON([("num", 1)])}, + ], + ) + + def test_gte(self): + with self.assertNumQueries(1) as ctx: + list(Number.objects.filter(num__gte=2)) + self.assertAggregateQuery( + ctx.captured_queries[0]["sql"], + "lookup__number", + [ + {"$match": {"num": {"$gte": 2}}}, + {"$addFields": {"num": "$num"}}, + {"$sort": SON([("num", 1)])}, + ], + ) + + def test_group_by_with_having(self): + with self.assertNumQueries(1) as ctx: + list(Number.objects.values("num").annotate(total=Sum("num")).filter(total=1)) + self.assertAggregateQuery( + ctx.captured_queries[0]["sql"], + "lookup__number", + [ + { + "$group": { + "__aggregation1": {"$sum": "$num"}, + "_id": {"num": "$num"}, + "total": {"$sum": "$num"}, + } + }, + {"$addFields": {"num": "$_id.num"}}, + {"$unset": "_id"}, + {"$match": {"__aggregation1": 1}}, + {"$project": {"num": 1, "total": "$__aggregation1"}}, + {"$sort": SON([("num", 1)])}, + ], + ) + + def test_subquery_filter_constant(self): + with self.assertNumQueries(1) as ctx: + list(Number.objects.filter(num__in=Number.objects.filter(num__gt=2).values("num"))) + self.assertAggregateQuery( + ctx.captured_queries[0]["sql"], + "lookup__number", + [ + { + "$lookup": { + "as": "__subquery0", + "from": "lookup__number", + "let": {}, + "pipeline": [ + {"$match": {"num": {"$gt": 2}}}, + { + "$facet": { + "group": [ + {"$group": {"_id": None, "tmp_name": {"$addToSet": "$num"}}} + ] + } + }, + { + "$project": { + "num": { + "$ifNull": [ + { + "$getField": { + "input": {"$arrayElemAt": ["$group", 0]}, + "field": "tmp_name", + } + }, + [], + ] + } + } + }, + ], + } + }, + { + "$set": { + "__subquery0": { + "$cond": { + "if": { + "$or": [ + {"$eq": [{"$type": "$__subquery0"}, "missing"]}, + {"$eq": [{"$size": "$__subquery0"}, 0]}, + ] + }, + "then": {}, + "else": {"$arrayElemAt": ["$__subquery0", 0]}, + } + } + } + }, + {"$match": {"$expr": {"$in": ["$num", "$__subquery0.num"]}}}, + {"$addFields": {"num": "$num"}}, + {"$sort": SON([("num", 1)])}, + ], + )