Unify QS handling for model/object permissions

Ryan P Kilby · Ryan P Kilby · commit e5eb3151b01e · 2017-09-01T13:37:58.000-04:00
diff --git a/rest_framework/permissions.py b/rest_framework/permissions.py
@@ -114,6 +114,17 @@ def get_required_permissions(self, method, model_cls):
 
         return [perm % kwargs for perm in self.perms_map[method]]
 
+    def _queryset(self, view):
+        assert hasattr(view, 'get_queryset') \
+            or getattr(view, 'queryset', None) is not None, (
+            'Cannot apply {} on a view that does not set '
+            '`.queryset` or have a `.get_queryset()` method.'
+        ).format(self.__class__.__name__)
+
+        if hasattr(view, 'get_queryset'):
+            return view.get_queryset()
+        return view.queryset
+
     def has_permission(self, request, view):
         # Workaround to ensure DjangoModelPermissions are not applied
         # to the root view when using DefaultRouter.
@@ -124,19 +135,7 @@ def has_permission(self, request, view):
            not is_authenticated(request.user) and self.authenticated_users_only):
             return False
 
-        if hasattr(view, 'get_queryset'):
-            queryset = view.get_queryset()
-            assert queryset is not None, (
-                '{}.get_queryset() returned None'.format(view.__class__.__name__)
-            )
-        else:
-            queryset = getattr(view, 'queryset', None)
-
-        assert queryset is not None, (
-            'Cannot apply DjangoModelPermissions on a view that '
-            'does not set `.queryset` or have a `.get_queryset()` method.'
-        )
-
+        queryset = self._queryset(view)
         perms = self.get_required_permissions(request.method, queryset.model)
 
         return request.user.has_perms(perms)
@@ -183,16 +182,8 @@ def get_required_object_permissions(self, method, model_cls):
         return [perm % kwargs for perm in self.perms_map[method]]
 
     def has_object_permission(self, request, view, obj):
-        if hasattr(view, 'get_queryset'):
-            queryset = view.get_queryset()
-        else:
-            queryset = getattr(view, 'queryset', None)
-
-        assert queryset is not None, (
-            'Cannot apply DjangoObjectPermissions on a view that '
-            'does not set `.queryset` or have a `.get_queryset()` method.'
-        )
-
+        # authentication checks have already executed via has_permission
+        queryset = self._queryset(view)
         model_cls = queryset.model
         user = request.user
 
diff --git a/tests/test_permissions.py b/tests/test_permissions.py
@@ -9,7 +9,7 @@
 
 from rest_framework import (
     HTTP_HEADER_ENCODING, authentication, generics, permissions, serializers,
-    status
+    status, views
 )
 from rest_framework.compat import ResolverMatch, guardian, set_many
 from rest_framework.filters import DjangoObjectPermissionsFilter
@@ -219,6 +219,27 @@ def get_queryset(_):
         response = view(request)
         self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED)
 
+    def test_queryset_assertions(self):
+        class View(views.APIView):
+            authentication_classes = [authentication.BasicAuthentication]
+            permission_classes = [permissions.DjangoModelPermissions]
+        view = View.as_view()
+
+        request = factory.get('/', HTTP_AUTHORIZATION=self.permitted_credentials)
+        msg = 'Cannot apply DjangoModelPermissions on a view that does not set `.queryset` or have a `.get_queryset()` method.'
+        with self.assertRaisesMessage(AssertionError, msg):
+            view(request)
+
+        # Faulty `get_queryset()` methods should trigger the above "view does not have a queryset" assertion.
+        class View(RootView):
+            def get_queryset(self):
+                return None
+        view = View.as_view()
+
+        request = factory.get('/', HTTP_AUTHORIZATION=self.permitted_credentials)
+        with self.assertRaisesMessage(AttributeError, "'NoneType' object has no attribute 'model'"):
+            view(request)
+
 
 class BasicPermModel(models.Model):
     text = models.CharField(max_length=100)
