Handle chunks and thumbnails parameters in StartProcess endpoint

arkindex/dataimport/api.py

@@ -28,7 +28,7 @@ from arkindex.dataimport.serializers.git import \
 from arkindex.dataimport.serializers.imports import (
     DataImportLightSerializer, DataImportSerializer,
     DataImportFromFilesSerializer, ElementsWorkflowSerializer, WorkerRunLightSerializer, WorkerRunSerializer,
-    ImportTranskribusSerializer
+    ImportTranskribusSerializer, StartProcessSerializer
 )
 from arkindex.dataimport.serializers.workers import WorkerSerializer, WorkerVersionSerializer, RepositorySerializer
 from arkindex.users.models import OAuthCredentials, User
@@ -383,7 +383,9 @@ class StartProcess(APIView):
         if dataimport.mode != DataImportMode.Workers or dataimport.workflow is not None:
             raise ValidationError('Only a DataImport with Workers mode and not already launched can be started later on')
 
-        dataimport.start()
+        serializer = StartProcessSerializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+        dataimport.start(**serializer.validated_data)
 
         return Response(
             status=status.HTTP_200_OK,

arkindex/dataimport/models.py

@@ -121,7 +121,7 @@ class DataImport(IndexableModel):
         elif self.mode in (DataImportMode.Elements, DataImportMode.Workers):
             if chunks is not None:
                 assert isinstance(chunks, int) and chunks > 0, 'Chunks count should be an positive integer'
-                assert chunks < 33, 'Import distribution is limited to 32 chunks'
+                assert chunks < 11, 'Import distribution is limited to 10 chunks'
                 ml_workflow_chunks = chunks
 
             args = [
@@ -211,7 +211,8 @@ class DataImport(IndexableModel):
 
             # Generate a task for each WorkerRun on the DataImport
             for worker_run in self.worker_runs.all():
-                tasks[f'{worker_run.version.slug}'] = worker_run.build_task_recipe(import_task_name, elements_path)
+                task_name = f'{worker_run.version.slug}{task_suffix}'
+                tasks[task_name] = worker_run.build_task_recipe(import_task_name, elements_path)
 
         # Build the workflow in db
         recipe = settings.PONOS_RECIPE.copy()

arkindex/dataimport/serializers/imports.py

@@ -171,6 +171,11 @@ class DataImportFromFilesSerializer(serializers.Serializer):
         return data
 
 
+class StartProcessSerializer(serializers.Serializer):
+    chunks = serializers.IntegerField(min_value=1, max_value=10, default=1)
+    thumbnails = serializers.BooleanField(default=False)
+
+
 class ElementsWorkflowSerializer(serializers.Serializer):
 
     ml_tools = MLToolTaskSerializer(many=True)

arkindex/dataimport/tests/test_imports.py

@@ -6,6 +6,8 @@ from arkindex.documents.models import Corpus, ElementType
 from arkindex.project.tests import FixtureAPITestCase
 from arkindex.users.models import User
 from ponos.models import State, Workflow
+from unittest.mock import patch, MagicMock
+import yaml
 import uuid
 
 RECIPE = '''
@@ -730,3 +732,43 @@ class TestImports(FixtureAPITestCase):
         dataimport2.refresh_from_db()
         self.assertEqual(dataimport2.state, State.Unscheduled)
         self.assertIsNotNone(dataimport2.workflow)
+
+    def test_start_process_params_validation(self):
+        dataimport = self.corpus.imports.create(creator=self.user, mode=DataImportMode.Workers)
+        self.client.force_login(self.user)
+        wrong_params_checks = [
+            ({'chunks': 0}, {'chunks': ['Ensure this value is greater than or equal to 1.']}),
+            ({'chunks': 20}, {'chunks': ['Ensure this value is less than or equal to 10.']}),
+            ({'chunks': 'max'}, {'chunks': ['A valid integer is required.']}),
+            ({'thumbnails': 'gloubiboulga'}, {'thumbnails': ['Must be a valid boolean.']})
+        ]
+        for (params, check) in wrong_params_checks:
+            response = self.client.post(reverse('api:process-start', kwargs={'pk': str(dataimport.id)}), params)
+            self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+            self.assertDictEqual(response.json(), check)
+
+    @patch('arkindex.dataimport.models.DataImport.worker_runs')
+    def test_start_process_parameters(self, worker_runs_mock):
+        """
+        It should be possible to pass chunks and thumbnails parameters when starting a workflow
+        """
+        dataimport = self.corpus.imports.create(creator=self.user, mode=DataImportMode.Workers)
+        # Add a worker run to this process
+        run_mock = MagicMock()
+        run_mock.version.slug = 'my_worker'
+        run_mock.build_task_recipe.return_value = {'image': ''}
+        worker_runs_mock.all.return_value = [run_mock]
+
+        self.client.force_login(self.user)
+        response = self.client.post(
+            reverse('api:process-start', kwargs={'pk': str(dataimport.id)}),
+            {'chunks': 3, 'thumbnails': 'true'}
+        )
+        self.assertEqual(response.status_code, status.HTTP_200_OK)
+        dataimport.refresh_from_db()
+        self.maxDiff = None
+        recipe_dump = yaml.safe_load(dataimport.workflow.recipe)
+        self.assertCountEqual(
+            recipe_dump['tasks'].keys(),
+            ['import', 'my_worker_1', 'my_worker_2', 'my_worker_3', 'thumbnails_1', 'thumbnails_2', 'thumbnails_3']
+        )