From c3876eb9c174017c2b55ce2d142ea1b652383d1b Mon Sep 17 00:00:00 2001
From: Jamie Howard <jamie.howard@strapi.io>
Date: Mon, 19 Jun 2023 11:27:45 +0100
Subject: [PATCH] chore: cleanup

---
 .../review-workflows/entity-service-decorator.js      | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/packages/core/admin/ee/server/services/review-workflows/entity-service-decorator.js b/packages/core/admin/ee/server/services/review-workflows/entity-service-decorator.js
index e15f7cf014..9ec3f292f8 100644
--- a/packages/core/admin/ee/server/services/review-workflows/entity-service-decorator.js
+++ b/packages/core/admin/ee/server/services/review-workflows/entity-service-decorator.js
@@ -18,6 +18,12 @@ const getDataWithStage = async (data) => {
   return data;
 };
 
+/**
+ * Get the stage information of an entity
+ * @param {String} uid
+ * @param {Number} id
+ * @returns {Object}
+ */
 const getEntityStage = async (uid, id) => {
   const entity = await strapi.entityService.findOne(uid, id, {
     populate: {
@@ -28,7 +34,7 @@ const getEntityStage = async (uid, id) => {
       },
     },
   });
-  return entity?.[ENTITY_STAGE_ATTRIBUTE] ?? null;
+  return entity?.[ENTITY_STAGE_ATTRIBUTE] ?? {};
 };
 
 /**
@@ -60,9 +66,8 @@ const decorator = (service) => ({
       return service.update.call(this, uid, entityId, { ...opts, data });
     }
 
-    const previousStage = await getEntityStage(uid, entityId);
-
     const updatedEntity = await service.update.call(this, uid, entityId, { ...opts, data });
+    const previousStage = await getEntityStage(uid, entityId);
     if (
       previousStage?.workflow?.id &&
       previousStage?.id &&