From d8da5b5214400bdcb102317ac9af011e2c12ed08 Mon Sep 17 00:00:00 2001
From: Jamie Howard <jamie.howard@strapi.io>
Date: Wed, 29 Mar 2023 14:43:47 +0100
Subject: [PATCH] fix(ee): logic for finding an entities target stage

---
 .../server/services/__tests__/stages.test.js  |  4 +-
 .../services/review-workflows/stages.js       | 66 +++++++++++++++----
 2 files changed, 55 insertions(+), 15 deletions(-)

diff --git a/packages/core/admin/ee/server/services/__tests__/stages.test.js b/packages/core/admin/ee/server/services/__tests__/stages.test.js
index 03d480d770..9e7e073422 100644
--- a/packages/core/admin/ee/server/services/__tests__/stages.test.js
+++ b/packages/core/admin/ee/server/services/__tests__/stages.test.js
@@ -156,7 +156,7 @@ describe('Review workflows - Stages service', () => {
     });
 
     test('Should move entities in a deleted stage to the previous stage', async () => {
-      await stagesService.replaceWorkflowStages(1, workflowMock.stages.slice(0, 1));
+      await stagesService.replaceWorkflowStages(1, workflowMock.stages.slice(0, 2));
 
       expect(servicesMock['admin::workflows'].findById).toBeCalled();
       expect(entityServiceMock.create).not.toBeCalled();
@@ -173,7 +173,7 @@ describe('Review workflows - Stages service', () => {
 
       expect(servicesMock['admin::workflows'].update).toBeCalled();
       expect(servicesMock['admin::workflows'].update).toBeCalledWith(workflowMock.id, {
-        stages: [workflowMock.stages[0].id],
+        stages: [workflowMock.stages[0].id, workflowMock.stages[1].id],
       });
     });
 
diff --git a/packages/core/admin/ee/server/services/review-workflows/stages.js b/packages/core/admin/ee/server/services/review-workflows/stages.js
index 31225fc3bb..d82b51ab29 100644
--- a/packages/core/admin/ee/server/services/review-workflows/stages.js
+++ b/packages/core/admin/ee/server/services/review-workflows/stages.js
@@ -65,6 +65,8 @@ module.exports = ({ strapi }) => {
         const defaultWorkflow = await getDefaultWorkflow({ strapi });
 
         await mapAsync(updated, (stage) => this.update(stage.id, stage));
+
+        const entitiesToMove = [];
         await mapAsync(deleted, async (stage) => {
           // Find any entities related to this stage
           const stageInfo = await this.findById(stage.id, {
@@ -73,25 +75,39 @@ module.exports = ({ strapi }) => {
           });
 
           if ((stageInfo?.related?.length ?? 0) === 0) {
+            // If there are no related entities, just delete the stage
             return this.delete(stage.id);
           }
 
-          // This stage has related entities that we need to move to the
-          // previous stage
-          const previousStage =
-            defaultWorkflow.stages[defaultWorkflow.stages.findIndex((s) => s.id === stage.id) - 1];
+          // This stage has related entities that need to be moved to their
+          // target stage
+          // Find the nearest stage in the workflow that is not deleted,
+          // prioritizing the previous stages
+          const targetStageId = findNearestMatchingStageID(
+            defaultWorkflow.stages,
+            defaultWorkflow.stages.findIndex((s) => s.id === stage.id),
+            (targetStage) => {
+              return !deleted.find((s) => s.id === targetStage.id);
+            }
+          );
+
+          // Keep track of the entities to move and their target stages
+          entitiesToMove.push(...stageInfo.related.map((entity) => ({ ...entity, targetStageId })));
 
-          mapAsync(stageInfo.related, (entity) => {
-            return this.updateEntity(
-              {
-                id: entity.id,
-                modelUID: entity.__type,
-              },
-              previousStage.id
-            );
-          });
           return this.delete(stage.id);
         });
+
+        // Move all the entities whose stage was deleted to their target stage
+        mapAsync(entitiesToMove, (entity) => {
+          return this.updateEntity(
+            {
+              id: entity.id,
+              modelUID: entity.__type,
+            },
+            entity.targetStageId
+          );
+        });
+
         return workflowsService.update(workflowId, {
           stages: stagesIds,
         });
@@ -174,3 +190,27 @@ function assertAtLeastOneStageRemain(workflowStages, diffStages) {
     throw new ApplicationError('At least one stage must remain in the workflow.');
   }
 }
+
+/**
+ * Find the nearest object in an array that matches a condition.
+ * Starts by searching the elements before the index, then the remaining elements in the array.
+ *
+ * @param {Array} array of stages
+ * @param {Number} index the index to start searching from
+ * @param {Function} condition must evaluate to true for the object to be considered a match
+ * @returns {Object}
+ */
+function findNearestMatchingStageID(array, index, condition) {
+  // Start by searching the elements before the index
+  for (let i = index; i >= 0; i -= 1) {
+    if (condition(array[i])) {
+      return array[i].id;
+    }
+  }
+
+  // If no matching element is found before the index,
+  // search the remaining elements in the array
+  const remainingArray = array.slice(index + 1);
+  const nearestObject = remainingArray.filter(condition)[0];
+  return nearestObject.id;
+}