Refactor aggregates and add coverage unit test

Author: WaVEV

django_mongodb_backend/aggregates.py

@@ -19,6 +19,7 @@
 
 def aggregate(self, compiler, connection, operator=None, resolve_inner_expression=False):
     agg_expression, *_ = self.get_source_expressions()
+    lhs_mql = None
     if self.filter is not None:
         try:
             lhs_mql = self.filter.as_mql(compiler, connection, as_expr=True)
@@ -29,12 +30,13 @@ def aggregate(self, compiler, connection, operator=None, resolve_inner_expressio
                 # Skip rows that don't meet the criteria.
                 default=Remove(),
             )
-            lhs_mql = agg_expression.as_mql(compiler, connection, as_expr=True)
         except FullResultSet:
-            lhs_mql = agg_expression.as_mql(compiler, connection, as_expr=True)
+            # TODO: CHECK UNIT TEST REACH HERE.
+            pass
         except EmptyResultSet:
-            lhs_mql = Value(None).as_mql(compiler, connection, as_expr=True)
-    else:
+            # TODO: CHECK UNIT TEST REACH HERE.
+            agg_expression = Remove()
+    if lhs_mql is None:
         lhs_mql = agg_expression.as_mql(compiler, connection, as_expr=True)
     if resolve_inner_expression:
         return lhs_mql
@@ -50,10 +52,11 @@ def count(self, compiler, connection, resolve_inner_expression=False):
     """
     agg_expression, *_ = self.get_source_expressions()
     if not self.distinct or resolve_inner_expression:
+        lhs_mql = None
         conditions = [IsNull(agg_expression, False)]
         if self.filter:
             try:
-                inner_expression = self.filter.as_mql(compiler, connection, as_expr=True)
+                lhs_mql = self.filter.as_mql(compiler, connection, as_expr=True)
             except NotSupportedError:
                 conditions.append(self.filter.condition)
                 condition = When(
@@ -71,7 +74,8 @@ def count(self, compiler, connection, resolve_inner_expression=False):
                 # Skip rows that don't meet the criteria.
                 default=Remove(),
             )
-        lhs_mql = inner_expression.as_mql(compiler, connection, as_expr=True)
+        if lhs_mql is None:
+            lhs_mql = inner_expression.as_mql(compiler, connection, as_expr=True)
         if resolve_inner_expression:
             return lhs_mql
         return {"$sum": lhs_mql}

tests/aggregation_/__init__.py

@@ -0,0 +1,25 @@
+from django.db import models
+
+
+class Author(models.Model):
+    name = models.CharField(max_length=100)
+    age = models.IntegerField()
+    friends = models.ManyToManyField("self", blank=True)
+    rating = models.FloatField(null=True)
+
+    def __str__(self):
+        return self.name
+
+
+class Book(models.Model):
+    isbn = models.CharField(max_length=9)
+    name = models.CharField(max_length=255)
+    pages = models.IntegerField()
+    rating = models.FloatField()
+    price = models.DecimalField(decimal_places=2, max_digits=6)
+    authors = models.ManyToManyField(Author)
+    contact = models.ForeignKey(Author, models.CASCADE, related_name="book_contact_set")
+    pubdate = models.DateField()
+
+    def __str__(self):
+        return self.name

tests/aggregation_/models.py

@@ -0,0 +1,79 @@
+import datetime
+from decimal import Decimal
+
+from django.db.models import (
+    Count,
+    Max,
+    Q,
+)
+from django.test import TestCase
+
+from .models import Author, Book
+
+
+class FilteredAggregateTests(TestCase):
+    @classmethod
+    def setUpTestData(cls):
+        cls.a1 = Author.objects.create(name="test", age=40)
+        cls.a2 = Author.objects.create(name="test2", age=60)
+        cls.a3 = Author.objects.create(name="test3", age=40)
+        cls.b1 = Book.objects.create(
+            isbn="159059725",
+            name="The Definitive Guide to Django: Web Development Done Right",
+            pages=447,
+            rating=4.5,
+            price=Decimal("30.00"),
+            contact=cls.a1,
+            pubdate=datetime.date(2007, 12, 6),
+        )
+        cls.b2 = Book.objects.create(
+            isbn="067232959",
+            name="Sams Teach Yourself Django in 24 Hours",
+            pages=528,
+            rating=3.0,
+            price=Decimal("30.00"),
+            contact=cls.a2,
+            pubdate=datetime.date(2008, 3, 3),
+        )
+        cls.b3 = Book.objects.create(
+            isbn="159059996",
+            name="Practical Django Projects",
+            pages=600,
+            rating=40.5,
+            price=Decimal("30.00"),
+            contact=cls.a3,
+            pubdate=datetime.date(2008, 6, 23),
+        )
+        cls.a1.friends.add(cls.a2)
+        cls.a1.friends.add(cls.a3)
+        cls.b1.authors.add(cls.a1)
+        cls.b1.authors.add(cls.a3)
+        cls.b2.authors.add(cls.a2)
+        cls.b3.authors.add(cls.a3)
+
+    def test_filtered_aggregate_empty_condition_distinct(self):
+        book = Book.objects.annotate(
+            ages=Count("authors__age", filter=Q(authors__in=[]), distinct=True),
+        ).get(pk=self.b1.pk)
+        self.assertEqual(book.ages, 0)
+        aggregate = Book.objects.aggregate(max_rating=Max("rating", filter=Q(rating__in=[])))
+        self.assertEqual(aggregate, {"max_rating": None})
+
+    def test_filtered_aggregate_full_condition(self):
+        book = Book.objects.annotate(
+            ages=Count(
+                "authors__age",
+                filter=~Q(pk__in=[]),
+            ),
+        ).get(pk=self.b1.pk)
+        self.assertEqual(book.ages, 2)
+        aggregate = Book.objects.aggregate(max_rating=Max("rating", filter=~Q(rating__in=[])))
+        self.assertEqual(aggregate, {"max_rating": 40.5})
+
+    def test_filtered_aggregate_full_condition_distinct(self):
+        book = Book.objects.annotate(
+            ages=Count("authors__age", filter=~Q(authors__in=[]), distinct=True),
+        ).get(pk=self.b1.pk)
+        self.assertEqual(book.ages, 1)
+        aggregate = Book.objects.aggregate(max_rating=Max("rating", filter=~Q(rating__in=[])))
+        self.assertEqual(aggregate, {"max_rating": 40.5})