Fixed #33651 -- Added support for prefetching GenericForeignKey.

Co-authored-by: revanthgss <revanthgss@almabase.com>
Co-authored-by: Mariusz Felisiak <felisiak.mariusz@gmail.com>

2 files changed, 69 insertions, 8 deletions

diff --git a/tests/contenttypes_tests/test_fields.py b/tests/contenttypes_tests/test_fields.py
index 418669140b..5510f34cd0 100644
--- a/tests/contenttypes_tests/test_fields.py
+++ b/tests/contenttypes_tests/test_fields.py
@@ -1,9 +1,11 @@
 import json
 
 from django.contrib.contenttypes.fields import GenericForeignKey
+from django.contrib.contenttypes.prefetch import GenericPrefetch
 from django.db import models
 from django.test import TestCase
 from django.test.utils import isolate_apps
+from django.utils.deprecation import RemovedInDjango60Warning
 
 from .models import Answer, Post, Question
 
@@ -22,14 +24,6 @@ class GenericForeignKeyTests(TestCase):
         ):
             Answer.question.get_content_type()
 
-    def test_incorrect_get_prefetch_queryset_arguments(self):
-        with self.assertRaisesMessage(
-            ValueError, "Custom queryset can't be used for this lookup."
-        ):
-            Answer.question.get_prefetch_queryset(
-                Answer.objects.all(), Answer.objects.all()
-            )
-
     def test_get_object_cache_respects_deleted_objects(self):
         question = Question.objects.create(text="Who?")
         post = Post.objects.create(title="Answer", parent=question)
@@ -59,3 +53,55 @@ class GenericRelationTests(TestCase):
         answer2 = Answer.objects.create(question=question)
         result = json.loads(Question.answer_set.field.value_to_string(question))
         self.assertCountEqual(result, [answer1.pk, answer2.pk])
+
+
+class GetPrefetchQuerySetDeprecation(TestCase):
+    def test_generic_relation_warning(self):
+        Question.objects.create(text="test")
+        questions = Question.objects.all()
+        msg = (
+            "get_prefetch_queryset() is deprecated. Use get_prefetch_querysets() "
+            "instead."
+        )
+        with self.assertWarnsMessage(RemovedInDjango60Warning, msg):
+            questions[0].answer_set.get_prefetch_queryset(questions)
+
+    def test_generic_foreign_key_warning(self):
+        answers = Answer.objects.all()
+        msg = (
+            "get_prefetch_queryset() is deprecated. Use get_prefetch_querysets() "
+            "instead."
+        )
+        with self.assertWarnsMessage(RemovedInDjango60Warning, msg):
+            Answer.question.get_prefetch_queryset(answers)
+
+
+class GetPrefetchQuerySetsTests(TestCase):
+    def test_duplicate_querysets(self):
+        question = Question.objects.create(text="What is your name?")
+        answer = Answer.objects.create(text="Joe", question=question)
+        answer = Answer.objects.get(pk=answer.pk)
+        msg = "Only one queryset is allowed for each content type."
+        with self.assertRaisesMessage(ValueError, msg):
+            models.prefetch_related_objects(
+                [answer],
+                GenericPrefetch(
+                    "question",
+                    [
+                        Question.objects.all(),
+                        Question.objects.filter(text__startswith="test"),
+                    ],
+                ),
+            )
+
+    def test_generic_relation_invalid_length(self):
+        Question.objects.create(text="test")
+        questions = Question.objects.all()
+        msg = (
+            "querysets argument of get_prefetch_querysets() should have a length of 1."
+        )
+        with self.assertRaisesMessage(ValueError, msg):
+            questions[0].answer_set.get_prefetch_querysets(
+                instances=questions,
+                querysets=[Answer.objects.all(), Question.objects.all()],
+            )
diff --git a/tests/contenttypes_tests/test_models.py b/tests/contenttypes_tests/test_models.py
index 95518bb3df..36c14cf56f 100644
--- a/tests/contenttypes_tests/test_models.py
+++ b/tests/contenttypes_tests/test_models.py
@@ -1,5 +1,6 @@
 from django.apps import apps
 from django.contrib.contenttypes.models import ContentType, ContentTypeManager
+from django.contrib.contenttypes.prefetch import GenericPrefetch
 from django.db import models
 from django.db.migrations.state import ProjectState
 from django.test import TestCase, override_settings
@@ -328,3 +329,17 @@ class ContentTypesMultidbTests(TestCase):
             1, using="other"
         ):
             ContentType.objects.get_for_model(Author)
+
+
+class GenericPrefetchTests(TestCase):
+    def test_values_queryset(self):
+        msg = "Prefetch querysets cannot use raw(), values(), and values_list()."
+        with self.assertRaisesMessage(ValueError, msg):
+            GenericPrefetch("question", [Author.objects.values("pk")])
+        with self.assertRaisesMessage(ValueError, msg):
+            GenericPrefetch("question", [Author.objects.values_list("pk")])
+
+    def test_raw_queryset(self):
+        msg = "Prefetch querysets cannot use raw(), values(), and values_list()."
+        with self.assertRaisesMessage(ValueError, msg):
+            GenericPrefetch("question", [Author.objects.raw("select pk from author")])