fix remaining tests

x-pack/filebeat/input/awss3/sqs_test.go

@@ -43,17 +43,9 @@ func TestSQSReceiver(t *testing.T) {
 		mockMsgHandler := NewMockSQSProcessor(ctrl)
 		msg := newSQSMessage(newS3Event("log.json"))
 
-		// Execute sqsReader and verify calls/state.
-		receiver := newSQSReaderInput(config{MaxNumberOfMessages: maxMessages}, aws.Config{})
-		receiver.log = logp.NewLogger(inputName)
-		receiver.sqs = mockSQS
-		receiver.msgHandler = mockMsgHandler
-		receiver.metrics = newInputMetrics("", nil, 0)
-		receiver.run(ctx)
-
 		// Initial ReceiveMessage for maxMessages.
 		mockSQS.EXPECT().
-			ReceiveMessage(gomock.Any(), gomock.Eq(maxMessages)).
+			ReceiveMessage(gomock.Any(), gomock.Any()).
 			Times(1).
 			DoAndReturn(func(_ context.Context, _ int) ([]types.Message, error) {
 				// Return single message.
@@ -83,6 +75,13 @@ func TestSQSReceiver(t *testing.T) {
 			Times(1).
 			Return(nil)
 
+		// Execute sqsReader and verify calls/state.
+		sqsReader := newSQSReaderInput(config{MaxNumberOfMessages: maxMessages}, aws.Config{})
+		sqsReader.log = logp.NewLogger(inputName)
+		sqsReader.sqs = mockSQS
+		sqsReader.msgHandler = mockMsgHandler
+		sqsReader.metrics = newInputMetrics("", nil, 0)
+		sqsReader.run(ctx)
 	})
 
 	t.Run("retry after ReceiveMessage error", func(t *testing.T) {
@@ -91,36 +90,39 @@ func TestSQSReceiver(t *testing.T) {
 
 		ctrl, ctx := gomock.WithContext(ctx, t)
 		defer ctrl.Finish()
-		mockAPI := NewMockSQSAPI(ctrl)
+		mockSQS := NewMockSQSAPI(ctrl)
 		mockMsgHandler := NewMockSQSProcessor(ctrl)
 
 		gomock.InOrder(
 			// Initial ReceiveMessage gets an error.
-			mockAPI.EXPECT().
-				ReceiveMessage(gomock.Any(), gomock.Eq(maxMessages)).
+			mockSQS.EXPECT().
+				ReceiveMessage(gomock.Any(), gomock.Any()).
 				Times(1).
 				DoAndReturn(func(_ context.Context, _ int) ([]types.Message, error) {
 					return nil, errFakeConnectivityFailure
 				}),
 			// After waiting for sqsRetryDelay, it retries.
-			mockAPI.EXPECT().
-				ReceiveMessage(gomock.Any(), gomock.Eq(maxMessages)).
+			mockSQS.EXPECT().
+				ReceiveMessage(gomock.Any(), gomock.Any()).
 				Times(1).
 				DoAndReturn(func(_ context.Context, _ int) ([]types.Message, error) {
 					cancel()
 					return nil, nil
 				}),
 		)
+		mockSQS.EXPECT().
+			GetQueueAttributes(gomock.Any(), gomock.Eq([]types.QueueAttributeName{sqsApproximateNumberOfMessages})).
+			DoAndReturn(func(_ context.Context, _ []types.QueueAttributeName) (map[string]string, error) {
+				return map[string]string{sqsApproximateNumberOfMessages: "10000"}, nil
+			}).AnyTimes()
 
-		// Execute SQSReceiver and verify calls/state.
-		receiver := &sqsReaderInput{
-			log:        logp.NewLogger(inputName),
-			config:     config{MaxNumberOfMessages: maxMessages},
-			sqs:        mockAPI,
-			msgHandler: mockMsgHandler,
-			metrics:    newInputMetrics("", nil, 0),
-		}
-		receiver.run(ctx)
+		// Execute SQSReader and verify calls/state.
+		sqsReader := newSQSReaderInput(config{MaxNumberOfMessages: maxMessages}, aws.Config{})
+		sqsReader.log = logp.NewLogger(inputName)
+		sqsReader.sqs = mockSQS
+		sqsReader.msgHandler = mockMsgHandler
+		sqsReader.metrics = newInputMetrics("", nil, 0)
+		sqsReader.run(ctx)
 	})
 }