Add handlerSpy support to SNS

packages/sns/lib/sns/AbstractSnsPublisherMonoSchema.ts

@@ -54,6 +54,7 @@ export abstract class AbstractSnsPublisherMonoSchema<MessagePayloadType extends
       } satisfies PublishCommandInput
       const command = new PublishCommand(input)
       await this.snsClient.send(command)
+      this.handleMessageProcessed(message, 'published')
     } catch (error) {
       this.handleError(error)
       throw error

packages/sns/lib/sns/AbstractSnsPublisherMultiSchema.ts

@@ -60,6 +60,7 @@ export abstract class AbstractSnsPublisherMultiSchema<MessagePayloadType extends
       } satisfies PublishCommandInput
       const command = new PublishCommand(input)
       await this.snsClient.send(command)
+      this.handleMessageProcessed(message, 'published')
     } catch (error) {
       this.handleError(error)
       throw error

packages/sns/lib/utils/snsMessageDeserializer.spec.ts

@@ -11,6 +11,7 @@ import { deserializeSNSMessage } from './snsMessageDeserializer'
 describe('messageDeserializer', () => {
   it('deserializes valid JSON', () => {
     const messagePayload: PERMISSIONS_MESSAGE_TYPE = {
+      id: '1',
       messageType: 'add',
       userIds: [1],
       permissions: ['perm'],

packages/sns/test/consumers/SnsSqsPermissionConsumerMonoSchema.ts

@@ -38,6 +38,7 @@ export class SnsSqsPermissionConsumerMonoSchema extends AbstractSnsSqsConsumerMo
   ) {
     super(dependencies, {
       messageSchema: PERMISSIONS_MESSAGE_SCHEMA,
+      handlerSpy: true,
       messageTypeField: 'messageType',
       deletionConfig: {
         deleteIfExists: true,

packages/sns/test/consumers/SnsSqsPermissionConsumerMultiSchema.ts

@@ -50,6 +50,7 @@ export class SnsSqsPermissionConsumerMultiSchema extends AbstractSnsSqsConsumerM
     super(
       dependencies,
       {
+        handlerSpy: true,
         handlers: new MessageHandlerConfigBuilder<SupportedEvents, ExecutionContext>()
           .addConfig(
             PERMISSIONS_ADD_MESSAGE_SCHEMA,

packages/sns/test/consumers/SnsSqsPermissionsConsumerMonoSchema.spec.ts

@@ -165,6 +165,7 @@ describe('SNS PermissionsConsumer', () => {
         userPermissionMap[300] = []
 
         await publisher.publish({
+          id: '1',
           messageType: 'add',
           userIds,
           permissions: perms,
@@ -186,6 +187,7 @@ describe('SNS PermissionsConsumer', () => {
         expect(users).toHaveLength(0)
 
         await publisher.publish({
+          id: '2',
           messageType: 'add',
           userIds,
           permissions: perms,
@@ -216,6 +218,7 @@ describe('SNS PermissionsConsumer', () => {
         userPermissionMap[100] = []
 
         await publisher.publish({
+          id: '3',
           messageType: 'add',
           userIds,
           permissions: perms,

packages/sns/test/consumers/SnsSqsPermissionsConsumerMultiSchema.spec.ts

@@ -1,6 +1,5 @@
 import type { SNSClient } from '@aws-sdk/client-sns'
 import type { SQSClient } from '@aws-sdk/client-sqs'
-import { waitAndRetry } from '@message-queue-toolkit/core'
 import { assertQueue } from '@message-queue-toolkit/sqs'
 import type { AwilixContainer } from 'awilix'
 import { describe, beforeEach, afterEach, expect, it, beforeAll } from 'vitest'
@@ -90,23 +89,21 @@ describe('SNS PermissionsConsumerMultiSchema', () => {
     describe('happy path', () => {
       it('Processes messages', async () => {
         await publisher.publish({
+          id: '1',
           messageType: 'add',
         })
         await publisher.publish({
+          id: '2',
           messageType: 'remove',
         })
         await publisher.publish({
+          id: '3',
           messageType: 'remove',
         })
 
-        await waitAndRetry(
-          () => {
-            return consumer.addCounter === 1 && consumer.removeCounter === 2
-          },
-          // Removing this makes test flaky
-          30,
-          20,
-        )
+        await consumer.handlerSpy.waitForMessageWithId('1', 'consumed')
+        await consumer.handlerSpy.waitForMessageWithId('2', 'consumed')
+        await consumer.handlerSpy.waitForMessageWithId('3', 'consumed')
 
         expect(consumer.addBarrierCounter).toBe(3)
         expect(consumer.addCounter).toBe(1)

packages/sns/test/consumers/userConsumerSchemas.ts

@@ -1,16 +1,19 @@
 import z from 'zod'
 
 export const PERMISSIONS_MESSAGE_SCHEMA = z.object({
+  id: z.string(),
   messageType: z.enum(['add', 'remove']),
   userIds: z.array(z.number()).describe('User IDs'),
   permissions: z.array(z.string()).nonempty().describe('List of user permissions'),
 })
 
 export const PERMISSIONS_ADD_MESSAGE_SCHEMA = z.object({
+  id: z.string(),
   messageType: z.literal('add'),
 })
 
 export const PERMISSIONS_REMOVE_MESSAGE_SCHEMA = z.object({
+  id: z.string(),
   messageType: z.literal('remove'),
 })
 

packages/sns/test/publishers/SnsPermissionPublisher.spec.ts

@@ -108,6 +108,7 @@ describe('SNSPermissionPublisher', () => {
       const { permissionPublisher } = diContainer.cradle
 
       const message = {
+        id: '1',
         userIds,
         messageType: 'add',
         permissions: perms,
@@ -158,6 +159,7 @@ describe('SNSPermissionPublisher', () => {
       })
 
       expect(receivedMessage).toEqual({
+        id: '1',
         messageType: 'add',
         permissions: ['perm1', 'perm2'],
         userIds: [100, 200, 300],

packages/sns/test/publishers/SnsPermissionPublisherMonoSchema.ts

@@ -24,6 +24,7 @@ export class SnsPermissionPublisherMonoSchema extends AbstractSnsPublisherMonoSc
       deletionConfig: {
         deleteIfExists: false,
       },
+      handlerSpy: true,
       messageSchema: PERMISSIONS_MESSAGE_SCHEMA,
       messageTypeField: 'messageType',
       ...options,

packages/sns/test/publishers/SnsPermissionPublisherMultiSchema.ts

@@ -33,6 +33,7 @@ export class SnsPermissionPublisherMultiSchema extends AbstractSnsPublisherMulti
         deleteIfExists: false,
       },
       messageSchemas: [PERMISSIONS_ADD_MESSAGE_SCHEMA, PERMISSIONS_REMOVE_MESSAGE_SCHEMA],
+      handlerSpy: true,
       messageTypeField: 'messageType',
       ...options,
     })

packages/sqs/lib/sqs/AbstractSqsConsumer.ts

@@ -189,8 +189,8 @@ export abstract class AbstractSqsConsumer<
         const deserializedMessage = this.deserializeMessage(message)
         if (deserializedMessage.error === 'abort') {
           await this.failProcessing(message)
-          const messageId = this.tryToExtractId(message)
 
+          const messageId = this.tryToExtractId(message)
           this.handleMessageProcessed(null, 'invalid_message', messageId.result)
           return
         }