From ed5fe3766d15cb1c869fe916731682f57039c39a Mon Sep 17 00:00:00 2001
From: Jamie Howard <jamie.howard@strapi.io>
Date: Fri, 9 Jun 2023 14:25:19 +0100
Subject: [PATCH] fix(ee): entity service decorator tests

---
 .../entity-service-decorator.test.js          | 19 +++++++++++++++----
 .../entity-service-decorator.js               |  5 +++--
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/packages/core/admin/ee/server/services/review-workflows/__tests__/entity-service-decorator.test.js b/packages/core/admin/ee/server/services/review-workflows/__tests__/entity-service-decorator.test.js
index 1b8d4b2bdf..07445a79ef 100644
--- a/packages/core/admin/ee/server/services/review-workflows/__tests__/entity-service-decorator.test.js
+++ b/packages/core/admin/ee/server/services/review-workflows/__tests__/entity-service-decorator.test.js
@@ -1,6 +1,7 @@
 'use strict';
 
 const { omit } = require('lodash/fp');
+const { WORKFLOW_UPDATE_STAGE } = require('../../../constants/webhookEvents');
 const { decorator } = require('../entity-service-decorator')();
 
 jest.mock('../../../utils');
@@ -80,6 +81,8 @@ describe('Entity service decorator', () => {
   });
 
   describe('Update', () => {
+    const uid = 'test-model';
+
     test('Calls original update for non review workflow content types', async () => {
       const entry = {
         id: 1,
@@ -105,15 +108,23 @@ describe('Entity service decorator', () => {
 
       const defaultService = {
         update: jest.fn(() => Promise.resolve(entry)),
+        findOne: jest.fn(() => {
+          return { strapi_reviewWorkflows_stage: { id: 2 } };
+        }),
+        emitEvent: jest.fn(),
       };
 
       const service = decorator(defaultService);
 
       const id = 1;
       const input = { data: { title: 'title ', strapi_reviewWorkflows_stage: 1 } };
-      await service.update('test-model', id, input);
+      await service.update(uid, id, input);
 
-      expect(defaultService.update).toHaveBeenCalledWith('test-model', id, {
+      expect(defaultService.emitEvent).toHaveBeenCalledWith(uid, WORKFLOW_UPDATE_STAGE, {
+        id,
+      });
+
+      expect(defaultService.update).toHaveBeenCalledWith(uid, id, {
         ...input,
         data: {
           ...input.data,
@@ -135,9 +146,9 @@ describe('Entity service decorator', () => {
 
       const id = 1;
       const input = { data: { title: 'title ', strapi_reviewWorkflows_stage: null } };
-      await service.update('test-model', id, input);
+      await service.update(uid, id, input);
 
-      expect(defaultService.update).toHaveBeenCalledWith('test-model', id, {
+      expect(defaultService.update).toHaveBeenCalledWith(uid, id, {
         ...input,
         data: {
           ...omit('strapi_reviewWorkflows_stage', input.data),
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 660cdaf760..b9fc75366a 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
@@ -44,15 +44,16 @@ const decorator = (service) => ({
     const data = { ...opts.data };
     if (isNull(data[ENTITY_STAGE_ATTRIBUTE])) {
       delete data[ENTITY_STAGE_ATTRIBUTE];
+      return service.update.call(this, uid, entityId, { ...opts, data });
     }
 
+    const updatedEntity = await service.update.call(this, uid, entityId, { ...opts, data });
+
     const entity = await service.findOne.call(this, uid, entityId, {
       populate: [ENTITY_STAGE_ATTRIBUTE],
     });
     const previousStageId = entity?.[ENTITY_STAGE_ATTRIBUTE]?.id ?? null;
 
-    const updatedEntity = await service.update.call(this, uid, entityId, { ...opts, data });
-
     if (previousStageId && previousStageId !== data[ENTITY_STAGE_ATTRIBUTE]) {
       await service.emitEvent.call(this, uid, WORKFLOW_UPDATE_STAGE, updatedEntity);
     }