KeyError when using CreateClassificationsSelection on a corpus ID that does not exist

arkindex/documents/api/ml.py

@@ -385,14 +385,15 @@ class ManageClassificationsSelection(SelectionMixin, CorpusACLMixin, CreateAPIVi
         serializer = self.get_serializer(data=request.data)
         serializer.is_valid(raise_exception=True)
 
-        corpus_id = serializer.validated_data['corpus_id']
-        corpus = self.get_corpus(corpus_id, role=Role.Contributor)
+        corpus = serializer.validated_data['corpus_id']
+        if not self.has_write_access(corpus):
+            raise PermissionDenied(detail='You do not have contributor access to this corpus.')
 
         mode = serializer.validated_data['mode']
         if mode == ClassificationMode.Create:
             return self.create(corpus, request, *args, **kwargs)
         elif mode == ClassificationMode.Validate:
-            elements = self.get_selection(corpus_id)
+            elements = self.get_selection(corpus.id)
             Classification.objects.filter(
                 element__in=elements,
                 state=ClassificationState.Pending,

arkindex/documents/serializers/ml.py

@@ -179,7 +179,7 @@ class ClassificationsSelectionSerializer(serializers.ModelSerializer):
     Serializer for the manual classification creation on the selection
     """
     ml_class = serializers.UUIDField(required=False)
-    corpus_id = serializers.UUIDField()
+    corpus_id = serializers.PrimaryKeyRelatedField(queryset=Corpus.objects.all())
     mode = EnumField(ClassificationMode)
 
     class Meta:
@@ -194,7 +194,7 @@ class ClassificationsSelectionSerializer(serializers.ModelSerializer):
         if not settings.ARKINDEX_FEATURES['selection']:
             raise ValidationError(['Selection is not available on this instance.'])
         data = super().validate(data)
-        corpus = Corpus.objects.get(id=data['corpus_id'])
+        corpus = data['corpus_id']
 
         if data['mode'] == ClassificationMode.Create:
             if 'ml_class' not in data:

arkindex/documents/tests/test_moderation.py

+import uuid
+
 from django.test import override_settings
 from django.urls import reverse
 from rest_framework import status
@@ -438,7 +440,7 @@ class TestClasses(FixtureAPITestCase):
     @override_settings(ARKINDEX_FEATURES={'selection': False})
     def test_classification_selection_no_selection(self):
         self.client.force_login(self.user)
-        with self.assertNumQueries(2):
+        with self.assertNumQueries(3):
             response = self.client.post(
                 reverse('api:classification-selection'),
                 data={'corpus_id': self.corpus.id, 'ml_class': self.text.id, 'mode': 'create'}
@@ -455,6 +457,36 @@ class TestClasses(FixtureAPITestCase):
             )
             self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
 
+    def test_classification_selection_writable_corpus(self):
+        self.corpus.memberships.filter(user=self.user).update(level=Role.Guest.value)
+
+        self.client.force_login(self.user)
+        with self.assertNumQueries(7):
+            response = self.client.post(
+                reverse('api:classification-selection'),
+                data={'corpus_id': self.corpus.id, 'ml_class': self.text.id, 'mode': 'create'}
+            )
+            self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
+
+        self.assertEqual(response.json(), {
+            'detail': 'You do not have contributor access to this corpus.',
+        })
+
+    def test_classification_selection_non_existing_corpus(self):
+        random_uuid = uuid.uuid4()
+        self.client.force_login(self.user)
+
+        with self.assertNumQueries(3):
+            response = self.client.post(
+                reverse('api:classification-selection'),
+                data={'corpus_id': random_uuid, 'ml_class': self.text.id, 'mode': 'create'}
+            )
+            self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+
+        self.assertEqual(response.json(), {
+            'corpus_id': [f'Invalid pk "{random_uuid}" - object does not exist.'],
+        })
+
     def test_classification_selection_wrong_corpus(self):
         line = MLClass.objects.create(name='line', corpus=self.private_corpus)
         self.client.force_login(self.user)
@@ -486,7 +518,7 @@ class TestClasses(FixtureAPITestCase):
 
     def test_classifications_selection_wrong_mode(self):
         self.client.force_login(self.user)
-        with self.assertNumQueries(2):
+        with self.assertNumQueries(3):
             response = self.client.post(
                 reverse('api:classification-selection'),
                 data={'corpus_id': self.corpus.id, 'ml_class': self.text.id, 'mode': 'test'}
@@ -529,7 +561,7 @@ class TestClasses(FixtureAPITestCase):
             [str(self.element.id), str(self.folder.id), str(act_x.id), str(act_y.id)],
         )
 
-        with self.assertNumQueries(10):
+        with self.assertNumQueries(9):
             response = self.client.post(
                 reverse('api:classification-selection'),
                 data={'corpus_id': self.corpus.id, 'ml_class': self.text.id, 'mode': 'create'}
@@ -588,7 +620,7 @@ class TestClasses(FixtureAPITestCase):
             format='json'
         )
 
-        with self.assertNumQueries(5):
+        with self.assertNumQueries(4):
             response = self.client.post(
                 reverse('api:classification-selection'),
                 data={'corpus_id': self.corpus.id, 'mode': 'validate'}