Add support for ArrayAgg aggregations (#183)

stphivos · Sep 6, 2024 · 26257c5 · 26257c5
1 parent 8d1047e
commit 26257c5
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 1 deletion.
diff --git a/django_mock_queries/constants.py b/django_mock_queries/constants.py
@@ -83,12 +83,14 @@
 AGGREGATES_MAX = 'MAX'
 AGGREGATES_MIN = 'MIN'
 AGGREGATES_AVG = 'AVG'
+AGGREGATES_ARRAY = 'ARRAY_AGG'
 AGGREGATES = (
     AGGREGATES_SUM,
     AGGREGATES_COUNT,
     AGGREGATES_MAX,
     AGGREGATES_MIN,
     AGGREGATES_AVG,
+    AGGREGATES_ARRAY,
 )
 
 DjangoQ = locate('django.db.models.Q')

diff --git a/django_mock_queries/query.py b/django_mock_queries/query.py
@@ -151,7 +151,8 @@ def aggregate(self, *args, **kwargs):
                     AGGREGATES_COUNT: lambda: len(values),
                     AGGREGATES_MAX: lambda: max(values),
                     AGGREGATES_MIN: lambda: min(values),
-                    AGGREGATES_AVG: lambda: sum(values) / len(values)
+                    AGGREGATES_AVG: lambda: sum(values) / len(values),
+                    AGGREGATES_ARRAY: lambda: values,
                 }[expr.function]()
 
             if len(values) == 0 and expr.function == AGGREGATES_COUNT:

diff --git a/tests/test_query.py b/tests/test_query.py
@@ -324,6 +324,19 @@ def test_query_aggregate_performs_avg_on_queryset_field(self):
             [x.foo for x in items if x.foo is not None]
         )
 
+    def test_query_aggregate_performs_array_on_queryset_field(self):
+        items = [
+            MockModel(foo=5),
+            MockModel(foo=10),
+            MockModel(foo=15),
+        ]
+        self.mock_set.add(*items)
+
+        expr = MagicMock(function=AGGREGATES_ARRAY, source_expressions=[MockModel(name='foo')])
+        result = self.mock_set.aggregate(expr)
+
+        assert result['foo__array_agg'] == [x.foo for x in items]
+
     def test_query_aggregate_with_none_only_field_values_performs_correct_aggregation(self):
         items = [
             MockModel(foo=None),