144 lines
5.7 KiB
Python
144 lines
5.7 KiB
Python
from django.db.models import Prefetch, prefetch_related_objects
|
|
from django.test import TestCase
|
|
|
|
from .models import Author, Book, Reader
|
|
|
|
|
|
class PrefetchRelatedObjectsTests(TestCase):
|
|
"""
|
|
Since prefetch_related_objects() is just the inner part of
|
|
prefetch_related(), only do basic tests to ensure its API hasn't changed.
|
|
"""
|
|
@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):
|
|
"""A m2m can be followed through another m2m."""
|
|
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']], # Charlotte - Poems, Jane Eyre
|
|
[['Amy']], # Anne - Poems
|
|
[['Amy'], []], # Emily - Poems, Wuthering Heights
|
|
[['Amy', 'Belinda']], # Jane - Sense and Sense
|
|
]
|
|
)
|
|
|
|
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])
|