Refs #32096 -- Added test for ArrayAgg over JSONField key transforms.

1 files changed, 35 insertions, 8 deletions

diff --git a/tests/postgres_tests/test_aggregates.py b/tests/postgres_tests/test_aggregates.py
index 19f9596576..b6cce4658b 100644
--- a/tests/postgres_tests/test_aggregates.py
+++ b/tests/postgres_tests/test_aggregates.py
@@ -1,6 +1,7 @@
 import json
 
 from django.db.models import CharField, F, OuterRef, Q, Subquery, Value
+from django.db.models.fields.json import KeyTransform
 from django.db.models.functions import Cast, Concat, Substr
 from django.test.utils import Approximate
 
@@ -20,10 +21,27 @@ except ImportError:
 class TestGeneralAggregate(PostgreSQLTestCase):
     @classmethod
     def setUpTestData(cls):
-        cls.agg1 = AggregateTestModel.objects.create(boolean_field=True, char_field='Foo1', integer_field=0)
-        AggregateTestModel.objects.create(boolean_field=False, char_field='Foo2', integer_field=1)
-        AggregateTestModel.objects.create(boolean_field=False, char_field='Foo4', integer_field=2)
-        AggregateTestModel.objects.create(boolean_field=True, char_field='Foo3', integer_field=0)
+        cls.aggs = AggregateTestModel.objects.bulk_create([
+            AggregateTestModel(boolean_field=True, char_field='Foo1', integer_field=0),
+            AggregateTestModel(
+                boolean_field=False,
+                char_field='Foo2',
+                integer_field=1,
+                json_field={'lang': 'pl'},
+            ),
+            AggregateTestModel(
+                boolean_field=False,
+                char_field='Foo4',
+                integer_field=2,
+                json_field={'lang': 'en'},
+            ),
+            AggregateTestModel(
+                boolean_field=True,
+                char_field='Foo3',
+                integer_field=0,
+                json_field={'breed': 'collie'},
+            ),
+        ])
 
     def test_array_agg_charfield(self):
         values = AggregateTestModel.objects.aggregate(arrayagg=ArrayAgg('char_field'))
@@ -79,6 +97,15 @@ class TestGeneralAggregate(PostgreSQLTestCase):
                 )
                 self.assertEqual(values, {'arrayagg': expected_output})
 
+    def test_array_agg_jsonfield(self):
+        values = AggregateTestModel.objects.aggregate(
+            arrayagg=ArrayAgg(
+                KeyTransform('lang', 'json_field'),
+                filter=Q(json_field__lang__isnull=False),
+            ),
+        )
+        self.assertEqual(values, {'arrayagg': ['pl', 'en']})
+
     def test_array_agg_filter(self):
         values = AggregateTestModel.objects.aggregate(
             arrayagg=ArrayAgg('integer_field', filter=Q(integer_field__gt=0)),
@@ -302,9 +329,9 @@ class TestGeneralAggregate(PostgreSQLTestCase):
 
     def test_string_agg_array_agg_filter_in_subquery(self):
         StatTestModel.objects.bulk_create([
-            StatTestModel(related_field=self.agg1, int1=0, int2=5),
-            StatTestModel(related_field=self.agg1, int1=1, int2=4),
-            StatTestModel(related_field=self.agg1, int1=2, int2=3),
+            StatTestModel(related_field=self.aggs[0], int1=0, int2=5),
+            StatTestModel(related_field=self.aggs[0], int1=1, int2=4),
+            StatTestModel(related_field=self.aggs[0], int1=2, int2=3),
         ])
         for aggregate, expected_result in (
             (
@@ -341,7 +368,7 @@ class TestGeneralAggregate(PostgreSQLTestCase):
         ).exclude(stringagg='').values('id')
         self.assertSequenceEqual(
             AggregateTestModel.objects.filter(id__in=Subquery(subquery)),
-            [self.agg1],
+            [self.aggs[0]],
         )