[B-1592] Call shutdown on handler when ctx done (#233)

Co-authored-by: Mirco Bordoni <mirco.bordoni@toolsforhumanity.com>

consumer.go

@@ -47,6 +47,7 @@ loop:
 		select {
 		case <-ctx.Done():
 			zap.S().Info("closing jobs channel")
+			c.handler.Shutdown()
 			close(jobs)
 			break loop
 		default:

consumer_test.go

@@ -38,6 +38,7 @@ type MsgHandler struct {
 	msgsReceivedCount     int
 	expectedMsg           TestMsg
 	expectedMsgAttributes interface{}
+	shutdownReceived      bool
 }
 
 func TestConsume(t *testing.T) {
@@ -109,7 +110,8 @@ func TestConsume_GracefulShutdown(t *testing.T) {
 		BatchSize:         batchSize,
 		ExtendEnabled:     true,
 	}
-	consumer := NewConsumer(awsCfg, config, &MsgHandler{})
+	msgHandler := MsgHandler{}
+	consumer := NewConsumer(awsCfg, config, &msgHandler)
 	go func() {
 		time.Sleep(time.Second * 1)
 		// Cancel context to trigger graceful shutdown
@@ -122,6 +124,11 @@ func TestConsume_GracefulShutdown(t *testing.T) {
 		os.Exit(1)
 	}()
 	consumer.Consume(ctx)
+
+	assert.Eventually(t, func() bool {
+		// Check that shutdown was called
+		return assert.Equal(t, true, msgHandler.shutdownReceived)
+	}, time.Second*2, time.Millisecond*100)
 }
 
 func createQueue(t *testing.T, ctx context.Context, awsCfg aws.Config, queueName string) *string {
@@ -164,6 +171,12 @@ func (m *MsgHandler) Run(ctx context.Context, msg *Message) error {
 	return err
 }
 
+func (m *MsgHandler) Shutdown() {
+	zap.S().Info("Shutting down")
+	m.shutdownReceived = true
+	// Do nothing
+}
+
 func sendTestMsg(t *testing.T, ctx context.Context, consumer *Consumer, queueUrl *string, expectedMsg TestMsg) TestMsg {
 	messageBodyBytes, err := json.Marshal(expectedMsg)
 	_, err = consumer.sqs.SendMessage(ctx, &sqs.SendMessageInput{

handler.go

@@ -4,4 +4,5 @@ import "context"
 
 type Handler interface {
 	Run(ctx context.Context, msg *Message) error
+	Shutdown()
 }