Fixed #34331 -- Added QuerySet.aiterator() support for prefetch_related().

1 files changed, 11 insertions, 7 deletions

diff --git a/tests/async/test_async_queryset.py b/tests/async/test_async_queryset.py
index e83157f5bd..374b4576f9 100644
--- a/tests/async/test_async_queryset.py
+++ b/tests/async/test_async_queryset.py
@@ -5,10 +5,10 @@ from datetime import datetime
 from asgiref.sync import async_to_sync, sync_to_async
 
 from django.db import NotSupportedError, connection
-from django.db.models import Sum
+from django.db.models import Prefetch, Sum
 from django.test import TestCase, skipIfDBFeature, skipUnlessDBFeature
 
-from .models import SimpleModel
+from .models import RelatedModel, SimpleModel
 
 
 class AsyncQuerySetTest(TestCase):
@@ -26,6 +26,9 @@ class AsyncQuerySetTest(TestCase):
             field=3,
             created=datetime(2022, 1, 1, 0, 0, 2),
         )
+        cls.r1 = RelatedModel.objects.create(simple=cls.s1)
+        cls.r2 = RelatedModel.objects.create(simple=cls.s2)
+        cls.r3 = RelatedModel.objects.create(simple=cls.s3)
 
     @staticmethod
     def _get_db_feature(connection_, feature_name):
@@ -48,11 +51,12 @@ class AsyncQuerySetTest(TestCase):
         self.assertCountEqual(results, [self.s1, self.s2, self.s3])
 
     async def test_aiterator_prefetch_related(self):
-        qs = SimpleModel.objects.prefetch_related("relatedmodels").aiterator()
-        msg = "Using QuerySet.aiterator() after prefetch_related() is not supported."
-        with self.assertRaisesMessage(NotSupportedError, msg):
-            async for m in qs:
-                pass
+        results = []
+        async for s in SimpleModel.objects.prefetch_related(
+            Prefetch("relatedmodel_set", to_attr="prefetched_relatedmodel")
+        ).aiterator():
+            results.append(s.prefetched_relatedmodel)
+        self.assertCountEqual(results, [[self.r1], [self.r2], [self.r3]])
 
     async def test_aiterator_invalid_chunk_size(self):
         msg = "Chunk size must be strictly positive."