Support narrowing spy type

packages/core/lib/queues/HandlerSpy.ts

@@ -82,8 +82,11 @@ export class HandlerSpy<MessagePayloadSchemas extends object> {
     )
   }
 
-  waitForMessageWithId(id: string, processingResult?: MessageProcessingResult) {
-    return this.waitForMessage(
+  waitForMessageWithId<T extends MessagePayloadSchemas>(
+    id: string,
+    processingResult?: MessageProcessingResult,
+  ) {
+    return this.waitForMessage<T>(
       // @ts-ignore
       {
         [this.messageIdField]: id,
@@ -92,24 +95,22 @@ export class HandlerSpy<MessagePayloadSchemas extends object> {
     )
   }
 
-  waitForMessage(
-    fields: DeepPartial<MessagePayloadSchemas>,
+  waitForMessage<T extends MessagePayloadSchemas>(
+    fields: DeepPartial<T>,
     processingResult?: MessageProcessingResult,
-  ): Promise<SpyResult<MessagePayloadSchemas>> {
+  ): Promise<SpyResult<T>> {
     const processedMessageEntry = Object.values(this.messageBuffer.items).find(
       // @ts-ignore
-      (spyResult: SpyResultCacheEntry<MessagePayloadSchemas>) => {
+      (spyResult: SpyResultCacheEntry<T>) => {
         return this.messageMatchesFilter(spyResult.value, fields, processingResult)
       },
     )
     if (processedMessageEntry) {
       return Promise.resolve(processedMessageEntry.value)
     }
 
-    let resolve: (
-      value: SpyResult<MessagePayloadSchemas> | PromiseLike<SpyResult<MessagePayloadSchemas>>,
-    ) => void
-    const spyPromise = new Promise<SpyResult<MessagePayloadSchemas>>((_resolve) => {
+    let resolve: (value: SpyResult<T> | PromiseLike<SpyResult<T>>) => void
+    const spyPromise = new Promise<SpyResult<T>>((_resolve) => {
       resolve = _resolve
     })
 

packages/core/test/queues/HandlerSpy.spec.ts

@@ -8,6 +8,12 @@ type Message = {
   payload?: Record<string, unknown>
 }
 
+type MessageB = {
+  id2: string
+  status2: string
+  payload?: Record<string, unknown>
+}
+
 type DeepMessage = {
   id: string
   status: string
@@ -24,11 +30,21 @@ const TEST_MESSAGE: Message = {
   status: 'done',
 }
 
+const TEST_MESSAGEB: MessageB = {
+  id2: 'abc',
+  status2: 'done',
+}
+
 const TEST_MESSAGE_2: Message = {
   id: 'abcd',
   status: 'inprogress',
 }
 
+const TEST_MESSAGE_2B: MessageB = {
+  id2: 'abcd',
+  status2: 'inprogress',
+}
+
 const TEST_MESSAGE_3: Message = {
   id: 'abcd',
   status: 'inprogress',
@@ -104,6 +120,26 @@ describe('HandlerSpy', () => {
       expect(message.message).toEqual(TEST_MESSAGE)
     })
 
+    it('Finds previously consumed event with type narrowing', async () => {
+      const spy = new HandlerSpy<Message | MessageB>()
+
+      spy.addProcessedMessage({
+        processingResult: 'consumed',
+        message: TEST_MESSAGE_2B,
+      })
+
+      spy.addProcessedMessage({
+        processingResult: 'consumed',
+        message: TEST_MESSAGEB,
+      })
+
+      const message = await spy.waitForMessage<MessageB>({
+        status2: 'done',
+      })
+
+      expect(message.message).toEqual(TEST_MESSAGEB)
+    })
+
     it('Finds previously consumed event with a deep match', async () => {
       const spy = new HandlerSpy<Message>()
 
@@ -272,6 +308,26 @@ describe('HandlerSpy', () => {
       expect(message.message).toEqual(TEST_MESSAGE)
     })
 
+    it('Finds previously consumed event with type narrowing', async () => {
+      const spy = new HandlerSpy<Message | MessageB>({
+        messageIdField: 'id2',
+      })
+
+      spy.addProcessedMessage({
+        processingResult: 'consumed',
+        message: TEST_MESSAGE_2B,
+      })
+
+      spy.addProcessedMessage({
+        processingResult: 'consumed',
+        message: TEST_MESSAGEB,
+      })
+
+      const message = await spy.waitForMessageWithId<MessageB>(TEST_MESSAGEB.id2)
+
+      expect(message.message).toEqual(TEST_MESSAGEB)
+    })
+
     it('Waits for an invalid message to be rejected by id', async () => {
       const spy = new HandlerSpy<Message>()