Addressed a crash when only prefetching base class references.

Objects prefetching tests were only covering the combined base and child references case.
charettes · Jul 12, 2019 · 78e5ea4 · 78e5ea4
1 parent cdb4c54
commit 78e5ea4
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/polymodels/models.py b/polymodels/models.py
@@ -36,7 +36,10 @@ def __call__(self, obj, with_prefetched_objects=False):
         if proxy:
             casted = copy_fields(casted, proxy)
         if with_prefetched_objects:
-            casted._prefetched_objects_cache.update(obj._prefetched_objects_cache)
+            try:
+                casted._prefetched_objects_cache.update(obj._prefetched_objects_cache)
+            except AttributeError:
+                casted._prefetched_objects_cache = obj._prefetched_objects_cache
         return casted
 
 

diff --git a/tests/test_managers.py b/tests/test_managers.py
@@ -108,6 +108,18 @@ def test_select_subclasses_prefetch_related(self):
         zoo.animals.add(animal, mammal, monkey)
         other_monkey = Monkey.objects.create(name='monkey')
         monkey.friends.add(other_monkey)
+        queryset = Animal.objects.select_subclasses().prefetch_related('zoos')
+        with self.assertNumQueries(2):
+            self.assertSequenceEqual(queryset, [
+                animal,
+                mammal,
+                monkey,
+                other_monkey,
+            ])
+            self.assertSequenceEqual(queryset[0].zoos.all(), [zoo])
+            self.assertSequenceEqual(queryset[1].zoos.all(), [zoo])
+            self.assertSequenceEqual(queryset[2].zoos.all(), [zoo])
+        # Test prefetch related combination.
         queryset = Animal.objects.select_subclasses().prefetch_related(
             'zoos',
             'mammal__monkey__friends',