[3.1.x] Refs #32096 -- Added test for ArrayAgg over JSONField key transforms.

Backport of 1d650ad019c1ab8e73d1e5b2587bb232c8ab35b6 from master

3 files changed, 47 insertions, 12 deletions

diff --git a/tests/postgres_tests/migrations/0002_create_test_models.py b/tests/postgres_tests/migrations/0002_create_test_models.py
index cb5f4c6d3e..e334057d22 100644
--- a/tests/postgres_tests/migrations/0002_create_test_models.py
+++ b/tests/postgres_tests/migrations/0002_create_test_models.py
@@ -202,7 +202,11 @@ class Migration(migrations.Migration):
                 ('boolean_field', models.BooleanField(null=True)),
                 ('char_field', models.CharField(max_length=30, blank=True)),
                 ('integer_field', models.IntegerField(null=True)),
-            ]
+                ('json_field', models.JSONField(null=True)),
+            ],
+            options={
+                'required_db_vendor': 'postgresql',
+            },
         ),
         migrations.CreateModel(
             name='StatTestModel',
@@ -215,7 +219,10 @@ class Migration(migrations.Migration):
                     models.SET_NULL,
                     null=True,
                 )),
-            ]
+            ],
+            options={
+                'required_db_vendor': 'postgresql',
+            },
         ),
         migrations.CreateModel(
             name='NowTestModel',
diff --git a/tests/postgres_tests/models.py b/tests/postgres_tests/models.py
index 464245fbab..a5bfc72fe7 100644
--- a/tests/postgres_tests/models.py
+++ b/tests/postgres_tests/models.py
@@ -154,16 +154,17 @@ class ArrayFieldSubclass(ArrayField):
         super().__init__(models.IntegerField())
 
 
-class AggregateTestModel(models.Model):
+class AggregateTestModel(PostgreSQLModel):
     """
     To test postgres-specific general aggregation functions
     """
     char_field = models.CharField(max_length=30, blank=True)
     integer_field = models.IntegerField(null=True)
     boolean_field = models.BooleanField(null=True)
+    json_field = models.JSONField(null=True)
 
 
-class StatTestModel(models.Model):
+class StatTestModel(PostgreSQLModel):
     """
     To test postgres-specific aggregation functions for statistics
     """
diff --git a/tests/postgres_tests/test_aggregates.py b/tests/postgres_tests/test_aggregates.py
index c97ed8e62d..8963ddbcf1 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)),
@@ -254,9 +281,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 (
             (
@@ -293,7 +320,7 @@ class TestGeneralAggregate(PostgreSQLTestCase):
         ).exclude(stringagg='').values('id')
         self.assertSequenceEqual(
             AggregateTestModel.objects.filter(id__in=Subquery(subquery)),
-            [self.agg1],
+            [self.aggs[0]],
         )