allow waitSubscribersExecution to wait for X trigger

Author: adrien2p

packages/medusa-test-utils/src/__tests__/events.spec.ts

@@ -1,5 +1,8 @@
 import { EventEmitter } from "events"
 import { waitSubscribersExecution } from "../events"
+import { setTimeout } from "timers/promises"
+
+jest.setTimeout(30000)
 
 // Mock the IEventBusModuleService
 class MockEventBus {
@@ -31,11 +34,12 @@ describe("waitSubscribersExecution", () => {
   describe("with no existing listeners", () => {
     it("should resolve when event is fired before timeout", async () => {
       const waitPromise = waitSubscribersExecution(TEST_EVENT, eventBus as any)
-      setTimeout(() => eventBus.emit(TEST_EVENT, "test-data"), 100).unref()
+      await setTimeout(100)
+      eventBus.emit(TEST_EVENT, "test-data")
 
       jest.advanceTimersByTime(100)
 
-      await expect(waitPromise).resolves.toEqual(["test-data"])
+      await expect(waitPromise).resolves.toEqual([["test-data"]])
     })
 
     it("should reject when timeout is reached before event is fired", async () => {
@@ -70,12 +74,29 @@ describe("waitSubscribersExecution", () => {
         `Timeout of ${customTimeout}ms exceeded while waiting for event "${TEST_EVENT}"`
       )
     })
+
+    it("should resolve when event is fired multiple times", async () => {
+      const waitPromise = waitSubscribersExecution(
+        TEST_EVENT,
+        eventBus as any,
+        { triggerCount: 2 }
+      )
+      eventBus.emit(TEST_EVENT, "test-data")
+      eventBus.emit(TEST_EVENT, "test-data")
+
+      const promisesRes = await waitPromise
+      const res = promisesRes.pop()
+      expect(res).toHaveLength(2)
+      expect(res[0]).toEqual(["test-data"])
+      expect(res[1]).toEqual(["test-data"])
+    })
   })
 
   describe("with existing listeners", () => {
     it("should resolve when all listeners complete successfully", async () => {
-      const listener = jest.fn().mockImplementation(() => {
-        return new Promise((resolve) => setTimeout(resolve, 200).unref())
+      const listener = jest.fn().mockImplementation(async () => {
+        await setTimeout(200)
+        return "res"
       })
 
       eventBus.eventEmitter_.on(TEST_EVENT, listener)
@@ -132,20 +153,49 @@ describe("waitSubscribersExecution", () => {
 
       expect(listener).not.toHaveBeenCalled()
     })
+
+    it("should resolve when event is fired multiple times", async () => {
+      const listener = jest.fn().mockImplementation(async () => {
+        await setTimeout(200)
+        return "res"
+      })
+
+      eventBus.eventEmitter_.on(TEST_EVENT, listener)
+
+      const waitPromise = waitSubscribersExecution(
+        TEST_EVENT,
+        eventBus as any,
+        {
+          triggerCount: 2,
+        }
+      )
+
+      eventBus.emit(TEST_EVENT, "test-data")
+      eventBus.emit(TEST_EVENT, "test-data")
+
+      const promisesRes = await waitPromise
+      const res = promisesRes.pop()
+      expect(res).toHaveLength(2)
+      expect(res[0]).toEqual("res")
+      expect(res[1]).toEqual("res")
+    })
   })
 
   describe("with multiple listeners", () => {
     it("should resolve when all listeners complete", async () => {
-      const listener1 = jest.fn().mockImplementation(() => {
-        return new Promise((resolve) => setTimeout(resolve, 100).unref())
+      const listener1 = jest.fn().mockImplementation(async () => {
+        await setTimeout(100)
+        return "res"
       })
 
-      const listener2 = jest.fn().mockImplementation(() => {
-        return new Promise((resolve) => setTimeout(resolve, 200).unref())
+      const listener2 = jest.fn().mockImplementation(async () => {
+        await setTimeout(200)
+        return "res"
       })
 
-      const listener3 = jest.fn().mockImplementation(() => {
-        return new Promise((resolve) => setTimeout(resolve, 300).unref())
+      const listener3 = jest.fn().mockImplementation(async () => {
+        await setTimeout(300)
+        return "res"
       })
 
       eventBus.eventEmitter_.on(TEST_EVENT, listener1)

packages/medusa-test-utils/src/events.ts

@@ -5,7 +5,10 @@ type EventBus = {
 }
 
 type WaitSubscribersExecutionOptions = {
+  /** Timeout in milliseconds for waiting for the event. Defaults to 15000ms. */
   timeout?: number
+  /** Number of times the event should be triggered before resolving. Defaults to 1. */
+  triggerCount?: number
 }
 
 // Map to hold pending promises for each event.
@@ -41,7 +44,7 @@ const createTimeoutPromise = (
 const doWaitSubscribersExecution = (
   eventName: string | symbol,
   eventBus: EventBus,
-  { timeout = 15000 }: WaitSubscribersExecutionOptions = {}
+  { timeout = 15000, triggerCount = 1 }: WaitSubscribersExecutionOptions = {}
 ): Promise<any> => {
   const eventEmitter = eventBus.eventEmitter_
   const subscriberPromises: Promise<any>[] = []
@@ -50,16 +53,25 @@ const doWaitSubscribersExecution = (
     eventName
   )
 
+  let currentCount = 0
+
   if (!eventEmitter.listeners(eventName).length) {
     let ok: (value?: any) => void
     const promise = new Promise((resolve) => {
       ok = resolve
     })
     subscriberPromises.push(promise)
 
+    let res: any[] = []
     const newListener = async (...args: any[]) => {
-      eventEmitter.removeListener(eventName, newListener)
-      ok(...args)
+      currentCount++
+      res.push(args)
+
+      if (currentCount >= triggerCount) {
+        eventEmitter.removeListener(eventName, newListener)
+        const res_ = triggerCount === 1 ? res[0] : res
+        ok(res_)
+      }
     }
 
     Object.defineProperty(newListener, "__isSubscribersExecutionWrapper", {
@@ -83,22 +95,38 @@ const doWaitSubscribersExecution = (
         nok = reject
       })
       subscriberPromises.push(promise)
+      let res: any[] = []
 
       const newListener = async (...args2: any[]) => {
-        // As soon as the subscriber is executed, we restore the original listener
-        eventEmitter.removeListener(eventName, newListener)
-        let listenerToAdd = listener
-        while (listenerToAdd.originalListener) {
-          listenerToAdd = listenerToAdd.originalListener
-        }
-        eventEmitter.on(eventName, listenerToAdd)
-
         try {
-          const res = await listener.apply(eventBus, args2)
-          ok(res)
+          const listenerRes = listener.apply(eventBus, args2)
+          if (typeof listenerRes?.then === "function") {
+            await listenerRes.then((res_) => {
+              res.push(res_)
+              currentCount++
+            })
+          } else {
+            res.push(listenerRes)
+            currentCount++
+          }
+
+          if (currentCount >= triggerCount) {
+            const res_ = triggerCount === 1 ? res[0] : res
+            ok(res_)
+          }
         } catch (error) {
           nok(error)
         }
+
+        if (currentCount >= triggerCount) {
+          // As soon as the subscriber is executed the required number of times, we restore the original listener
+          eventEmitter.removeListener(eventName, newListener)
+          let listenerToAdd = listener
+          while (listenerToAdd.originalListener) {
+            listenerToAdd = listenerToAdd.originalListener
+          }
+          eventEmitter.on(eventName, listenerToAdd)
+        }
       }
 
       Object.defineProperty(newListener, "__isSubscribersExecutionWrapper", {
@@ -130,7 +158,7 @@ const doWaitSubscribersExecution = (
  *
  * @param eventName - The name of the event to wait for.
  * @param eventBus - The event bus instance.
- * @param options - Options including timeout.
+ * @param options - Options including timeout and triggerCount.
  */
 export const waitSubscribersExecution = (
   eventName: string | symbol,