Skip to content

Commit

Permalink
Add batch sqs ack unit test
Browse files Browse the repository at this point in the history
  • Loading branch information
Jeffail committed Apr 23, 2024
1 parent 5ab5b6f commit 4561daf
Showing 1 changed file with 84 additions and 0 deletions.
84 changes: 84 additions & 0 deletions internal/impl/aws/input_sqs_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package aws

import (
"context"
"fmt"
"strconv"
"testing"
"time"
Expand All @@ -11,6 +12,8 @@ import (
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/aws/aws-sdk-go-v2/service/sqs"
"github.com/aws/aws-sdk-go-v2/service/sqs/types"
"github.com/benthosdev/benthos/v4/public/service"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

Expand Down Expand Up @@ -81,6 +84,8 @@ func (m *mockSqsInput) ChangeMessageVisibilityBatch(ctx context.Context, input *
for _, entry := range input.Entries {
if _, found := m.mesTimeouts[*entry.Id]; found {
m.mesTimeouts[*entry.Id] = entry.VisibilityTimeout
} else {
panic("nope")
}
}

Expand Down Expand Up @@ -200,3 +205,82 @@ func TestSQSInput(t *testing.T) {
return msgsLen == 0
}, 5*time.Second, time.Second)
}

func TestSQSInputBatchAck(t *testing.T) {
tCtx := context.Background()
defer tCtx.Done()

messages := []types.Message{}
for i := 0; i < 101; i++ {
messages = append(messages, types.Message{
Body: aws.String(fmt.Sprintf("message-%v", i)),
MessageId: aws.String(fmt.Sprintf("id-%v", i)),
ReceiptHandle: aws.String(fmt.Sprintf("h-%v", i)),
})
}
expectedMessages := len(messages)

conf, err := config.LoadDefaultConfig(context.Background(),
config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider("xxxxx", "xxxxx", "xxxxx")),
)
require.NoError(t, err)

r, err := newAWSSQSReader(
sqsiConfig{
URL: "http://foo.example.com",
WaitTimeSeconds: 0,
DeleteMessage: true,
ResetVisibility: true,
MaxNumberOfMessages: 10,
},
conf,
nil,
)
require.NoError(t, err)

mockInput := &mockSqsInput{
mtx: make(chan struct{}, 1),
queueTimeout: 10,
messages: messages,
mesTimeouts: make(map[string]int32, expectedMessages),
}
mockInput.mtx <- struct{}{}
r.sqs = mockInput
go mockInput.TimeoutLoop(tCtx)

defer r.closeSignal.TriggerHardStop()
err = r.Connect(tCtx)
require.NoError(t, err)

receivedMessageAcks := map[string]service.AckFunc{}

for _, eMsg := range messages {
m, aFn, err := r.Read(tCtx)
require.NoError(t, err)

mBytes, err := m.AsBytes()
require.NoError(t, err)

assert.Equal(t, *eMsg.Body, string(mBytes))
receivedMessageAcks[string(mBytes)] = aFn
}

// Check that messages haven't been deleted from the queue
mockInput.do(func() {
require.Len(t, mockInput.messages, expectedMessages)
require.Len(t, mockInput.mesTimeouts, expectedMessages)
})

// Ack all messages as a batch
for _, aFn := range receivedMessageAcks {
require.NoError(t, aFn(tCtx, err))
}

require.Eventually(t, func() bool {
msgsLen := 0
mockInput.do(func() {
msgsLen = len(mockInput.messages)
})
return msgsLen == 0
}, 5*time.Second, time.Second)
}

0 comments on commit 4561daf

Please sign in to comment.