From 1d650ad019c1ab8e73d1e5b2587bb232c8ab35b6 Mon Sep 17 00:00:00 2001
From: Mariusz Felisiak <felisiak.mariusz@gmail.com>
Date: Mon, 12 Oct 2020 12:52:58 +0200
Subject: [PATCH] Refs #32096 -- Added test for ArrayAgg over JSONField key
 transforms.

---
 .../migrations/0002_create_test_models.py     | 11 ++++-
 tests/postgres_tests/models.py                |  5 ++-
 tests/postgres_tests/test_aggregates.py       | 43 +++++++++++++++----
 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 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]],
         )