diff options
Diffstat (limited to 'tests/postgres_tests/test_aggregates.py')
| -rw-r--r-- | tests/postgres_tests/test_aggregates.py | 43 |
1 files changed, 43 insertions, 0 deletions
diff --git a/tests/postgres_tests/test_aggregates.py b/tests/postgres_tests/test_aggregates.py index b5474d361e..e781fd87ee 100644 --- a/tests/postgres_tests/test_aggregates.py +++ b/tests/postgres_tests/test_aggregates.py @@ -365,6 +365,49 @@ class TestGeneralAggregate(PostgreSQLTestCase): ) self.assertCountEqual(qs.get(), [1, 2]) + def test_array_agg_filter_index(self): + aggr1 = AggregateTestModel.objects.create(integer_field=1) + aggr2 = AggregateTestModel.objects.create(integer_field=2) + StatTestModel.objects.bulk_create( + [ + StatTestModel(related_field=aggr1, int1=1, int2=0), + StatTestModel(related_field=aggr1, int1=2, int2=1), + StatTestModel(related_field=aggr2, int1=3, int2=0), + StatTestModel(related_field=aggr2, int1=4, int2=1), + ] + ) + qs = ( + AggregateTestModel.objects.filter(pk__in=[aggr1.pk, aggr2.pk]) + .annotate( + array=ArrayAgg("stattestmodel__int1", filter=Q(stattestmodel__int2=0)) + ) + .annotate(array_value=F("array__0")) + .values_list("array_value", flat=True) + ) + self.assertCountEqual(qs, [1, 3]) + + def test_array_agg_filter_slice(self): + aggr1 = AggregateTestModel.objects.create(integer_field=1) + aggr2 = AggregateTestModel.objects.create(integer_field=2) + StatTestModel.objects.bulk_create( + [ + StatTestModel(related_field=aggr1, int1=1, int2=0), + StatTestModel(related_field=aggr1, int1=2, int2=1), + StatTestModel(related_field=aggr2, int1=3, int2=0), + StatTestModel(related_field=aggr2, int1=4, int2=1), + StatTestModel(related_field=aggr2, int1=5, int2=0), + ] + ) + qs = ( + AggregateTestModel.objects.filter(pk__in=[aggr1.pk, aggr2.pk]) + .annotate( + array=ArrayAgg("stattestmodel__int1", filter=Q(stattestmodel__int2=0)) + ) + .annotate(array_value=F("array__1_2")) + .values_list("array_value", flat=True) + ) + self.assertCountEqual(qs, [[], [5]]) + def test_bit_and_general(self): values = AggregateTestModel.objects.filter(integer_field__in=[0, 1]).aggregate( bitand=BitAnd("integer_field") |