since prefetchrelatedobjects is just the inner part of prefetchrelated only do basic tests to ensure its api hasn t changed a m2m can be followed through another m2m s list objects all with self assertnumqueries2 prefetchrelatedobjectss booksreadby with self assertnumqueries0 self assertequal strr for r in b readby all for b in a books all for a in s amy belinda charlotte poems jane eyre amy anne poems amy emily poems wuthering heights amy belinda jane sense and sense def testprefetchobjectself book1 book objects getidself book1 id with self assertnumqueries1 prefetchrelatedobjectsbook1 prefetchs with self assertnumqueries0 self assertcountequal book1 s all self 1 self 2 self 3 def testprefetchobjecttwiceself book1 book objects getidself book1 id book2 book objects getidself book2 id with self assertnumqueries1 prefetchrelatedobjectsbook1 prefetchs with self assertnumqueries1 prefetchrelatedobjectsbook1 book2 prefetchs with self assertnumqueries0 self assertcountequalbook2 s all self 1 def testprefetchobjecttoattrself book1 book objects getidself book1 id with self assertnumqueries1 prefetchrelatedobjects book1 prefetchs toattrthes with self assertnumqueries0 self assertcountequal book1 thes self 1 self 2 self 3 def testprefetchobjecttoattrtwiceself book1 book objects getidself book1 id book2 book objects getidself book2 id with self assertnumqueries1 prefetchrelatedobjects book1 prefetchs toattrthes with self assertnumqueries1 prefetchrelatedobjects book1 book2 prefetchs toattrthes with self assertnumqueries0 self assertcountequalbook2 thes self 1 def testprefetchquerysetself book1 book objects getidself book1 id with self assertnumqueries1 prefetchrelatedobjects book1 prefetch s queryset objects filter idinself 1 id self 2 id with self assertnumqueries0 self assertcountequalbook1 s all self 1 self 2 since prefetch_related_objects is just the inner part of prefetch_related only do basic tests to ensure its api hasn t changed a m2m can be followed through another m2m charlotte poems jane eyre anne poems emily poems wuthering heights jane sense and sense
from django.db.models import Prefetch, prefetch_related_objects
from django.test import TestCase

from .models import Author, Book, Reader


class PrefetchRelatedObjectsTests(TestCase):
    

    @classmethod
    def setUpTestData(cls):
        cls.book1 = Book.objects.create(title="Poems")
        cls.book2 = Book.objects.create(title="Jane Eyre")
        cls.book3 = Book.objects.create(title="Wuthering Heights")
        cls.book4 = Book.objects.create(title="Sense and Sensibility")

        cls.author1 = Author.objects.create(name="Charlotte", first_book=cls.book1)
        cls.author2 = Author.objects.create(name="Anne", first_book=cls.book1)
        cls.author3 = Author.objects.create(name="Emily", first_book=cls.book1)
        cls.author4 = Author.objects.create(name="Jane", first_book=cls.book4)

        cls.book1.authors.add(cls.author1, cls.author2, cls.author3)
        cls.book2.authors.add(cls.author1)
        cls.book3.authors.add(cls.author3)
        cls.book4.authors.add(cls.author4)

        cls.reader1 = Reader.objects.create(name="Amy")
        cls.reader2 = Reader.objects.create(name="Belinda")

        cls.reader1.books_read.add(cls.book1, cls.book4)
        cls.reader2.books_read.add(cls.book2, cls.book4)

    def test_unknown(self):
        book1 = Book.objects.get(id=self.book1.id)
        with self.assertRaises(AttributeError):
            prefetch_related_objects([book1], "unknown_attribute")

    def test_m2m_forward(self):
        book1 = Book.objects.get(id=self.book1.id)
        with self.assertNumQueries(1):
            prefetch_related_objects([book1], "authors")

        with self.assertNumQueries(0):
            self.assertCountEqual(
                book1.authors.all(), [self.author1, self.author2, self.author3]
            )

    def test_m2m_reverse(self):
        author1 = Author.objects.get(id=self.author1.id)
        with self.assertNumQueries(1):
            prefetch_related_objects([author1], "books")

        with self.assertNumQueries(0):
            self.assertCountEqual(author1.books.all(), [self.book1, self.book2])

    def test_foreignkey_forward(self):
        authors = list(Author.objects.all())
        with self.assertNumQueries(1):
            prefetch_related_objects(authors, "first_book")

        with self.assertNumQueries(0):
            [author.first_book for author in authors]

    def test_foreignkey_reverse(self):
        books = list(Book.objects.all())
        with self.assertNumQueries(1):
            prefetch_related_objects(books, "first_time_authors")

        with self.assertNumQueries(0):
            [list(book.first_time_authors.all()) for book in books]

    def test_m2m_then_m2m(self):
        
        authors = list(Author.objects.all())
        with self.assertNumQueries(2):
            prefetch_related_objects(authors, "books__read_by")

        with self.assertNumQueries(0):
            self.assertEqual(
                [
                    [[str(r) for r in b.read_by.all()] for b in a.books.all()]
                    for a in authors
                ],
                [
                    [["Amy"], ["Belinda"]],  
                    [["Amy"]],  
                    [["Amy"], []],  
                    [["Amy", "Belinda"]],  
                ],
            )

    def test_prefetch_object(self):
        book1 = Book.objects.get(id=self.book1.id)
        with self.assertNumQueries(1):
            prefetch_related_objects([book1], Prefetch("authors"))

        with self.assertNumQueries(0):
            self.assertCountEqual(
                book1.authors.all(), [self.author1, self.author2, self.author3]
            )

    def test_prefetch_object_twice(self):
        book1 = Book.objects.get(id=self.book1.id)
        book2 = Book.objects.get(id=self.book2.id)
        with self.assertNumQueries(1):
            prefetch_related_objects([book1], Prefetch("authors"))
        with self.assertNumQueries(1):
            prefetch_related_objects([book1, book2], Prefetch("authors"))
        with self.assertNumQueries(0):
            self.assertCountEqual(book2.authors.all(), [self.author1])

    def test_prefetch_object_to_attr(self):
        book1 = Book.objects.get(id=self.book1.id)
        with self.assertNumQueries(1):
            prefetch_related_objects(
                [book1], Prefetch("authors", to_attr="the_authors")
            )

        with self.assertNumQueries(0):
            self.assertCountEqual(
                book1.the_authors, [self.author1, self.author2, self.author3]
            )

    def test_prefetch_object_to_attr_twice(self):
        book1 = Book.objects.get(id=self.book1.id)
        book2 = Book.objects.get(id=self.book2.id)
        with self.assertNumQueries(1):
            prefetch_related_objects(
                [book1],
                Prefetch("authors", to_attr="the_authors"),
            )
        with self.assertNumQueries(1):
            prefetch_related_objects(
                [book1, book2],
                Prefetch("authors", to_attr="the_authors"),
            )
        with self.assertNumQueries(0):
            self.assertCountEqual(book2.the_authors, [self.author1])

    def test_prefetch_queryset(self):
        book1 = Book.objects.get(id=self.book1.id)
        with self.assertNumQueries(1):
            prefetch_related_objects(
                [book1],
                Prefetch(
                    "authors",
                    queryset=Author.objects.filter(
                        id__in=[self.author1.id, self.author2.id]
                    ),
                ),
            )

        with self.assertNumQueries(0):
            self.assertCountEqual(book1.authors.all(), [self.author1, self.author2])