diff --git a/django_mock_queries/constants.py b/django_mock_queries/constants.py index d857ba1..9f67e1e 100644 --- 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 index d58f1ed..93d2ced 100644 --- 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 index 2b37d42..c11f4e5 100644 --- 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),