Initial Test cases setup and some refactorings

Author: Pankaj Agrawal

powertools-sqs/src/main/java/software/amazon/lambda/powertools/sqs/PowertoolsSqs.java

@@ -13,6 +13,8 @@
  */
 package software.amazon.lambda.powertools.sqs;
 
+import java.lang.reflect.Constructor;
+import java.lang.reflect.InvocationTargetException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Function;
@@ -100,8 +102,16 @@ public static <R> List<R> partialBatchProcessor(final SQSEvent event,
                                                     final Class<? extends SqsMessageHandler<R>> handler) {
 
         try {
-            return partialBatchProcessor(event, suppressException, handler.newInstance());
-        } catch (IllegalAccessException | InstantiationException e) {
+            SqsMessageHandler<R> handlerInstance;
+            if (null == handler.getDeclaringClass()) {
+                handlerInstance = handler.newInstance();
+            } else {
+                Constructor<? extends SqsMessageHandler<R>> constructor = handler.getDeclaredConstructor(handler.getDeclaringClass());
+                constructor.setAccessible(true);
+                handlerInstance = constructor.newInstance(handler.getDeclaringClass().newInstance());
+            }
+            return partialBatchProcessor(event, suppressException, handlerInstance);
+        } catch (IllegalAccessException | InstantiationException | NoSuchMethodException | InvocationTargetException e) {
             LOG.error("Failed invoking process method on handler", e);
             throw new RuntimeException("Unexpected error occurred. Please raise issue at " +
                     "https://github.com/awslabs/aws-lambda-powertools-java/issues", e);
@@ -124,7 +134,12 @@ public static <R> List<R> partialBatchProcessor(final SQSEvent event,
             }
         }
 
-        batchContext.processSuccessAndReset(suppressException);
+        try {
+            batchContext.processSuccessAndReset(suppressException);
+        } catch (SQSBatchProcessingException e) {
+            e.addSuccessMessageReturnValues(handlerReturn);
+            throw e;
+        }
 
         return handlerReturn;
     }

powertools-sqs/src/main/java/software/amazon/lambda/powertools/sqs/SQSBatchProcessingException.java

@@ -3,20 +3,56 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import com.amazonaws.services.lambda.runtime.events.SQSEvent;
+
 import static java.util.stream.Collectors.joining;
 
 public class SQSBatchProcessingException extends RuntimeException {
 
     private final List<Exception> exceptions;
+    private final List<SQSEvent.SQSMessage> failures;
+    private final List<Object> returnValues;
 
-    public SQSBatchProcessingException(List<Exception> exceptions) {
+    public <T> SQSBatchProcessingException(final List<Exception> exceptions,
+                                           final List<SQSEvent.SQSMessage> failures,
+                                           final List<T> successReturns) {
         this.exceptions = new ArrayList<>(exceptions);
+        this.failures = new ArrayList<>(failures);
+        this.returnValues = new ArrayList<>(successReturns);
+    }
+
+    public List<Exception> getExceptions() {
+        return exceptions;
+    }
+
+    public List<Object> successMessageReturnValues() {
+        return returnValues;
+    }
+
+    public List<SQSEvent.SQSMessage> getFailures() {
+        return failures;
     }
 
     @Override
     public String getMessage() {
         return exceptions.stream()
-                .map(Throwable::getMessage)
+                .map(Throwable::toString)
                 .collect(joining("\n"));
     }
+
+    @Override
+    public void printStackTrace() {
+        for (Exception exception : exceptions) {
+            exception.printStackTrace();
+        }
+    }
+
+    @Override
+    public String toString() {
+        return getMessage();
+    }
+
+    <T> void addSuccessMessageReturnValues(final List<T> returnValues) {
+        this.returnValues.addAll(returnValues);
+    }
 }

powertools-sqs/src/main/java/software/amazon/lambda/powertools/sqs/internal/BatchContext.java

@@ -1,5 +1,6 @@
 package software.amazon.lambda.powertools.sqs.internal;
 
+import java.awt.geom.Area;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -36,7 +37,7 @@ public void addFailure(SQSMessage event, Exception e) {
         exceptions.add(e);
     }
 
-    public void processSuccessAndReset(final boolean suppressException) {
+    public <T> void processSuccessAndReset(final boolean suppressException) {
         try {
             if (hasFailures()) {
 
@@ -46,7 +47,7 @@ public void processSuccessAndReset(final boolean suppressException) {
                     List<String> messageIds = failures.stream().map(SQSMessage::getMessageId).collect(toList());
                     LOG.debug(format("[%s] records failed processing, but exceptions are suppressed. Failed messages %s", failures.size(), messageIds));
                 } else {
-                    throw new SQSBatchProcessingException(exceptions);
+                    throw new SQSBatchProcessingException(exceptions, failures, new ArrayList<T>());
                 }
             }
         } finally {

powertools-sqs/src/main/java/software/amazon/lambda/powertools/sqs/internal/SqsMessageBatchProcessorAspect.java

@@ -29,7 +29,7 @@ && placedOnSqsEventRequestHandler(pjp)) {
 
             SQSEvent sqsEvent = (SQSEvent) proceedArgs[0];
 
-            partialBatchProcessor(sqsEvent, sqsBatchProcessor.suppressException(), sqsBatchProcessor.value().newInstance());
+            partialBatchProcessor(sqsEvent, sqsBatchProcessor.suppressException(), sqsBatchProcessor.value());
         }
 
         return pjp.proceed(proceedArgs);

powertools-sqs/src/test/java/software/amazon/lambda/powertools/sqs/PowertoolsSqsBatchProcessorTest.java

@@ -0,0 +1,76 @@
+package software.amazon.lambda.powertools.sqs;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.stream.Stream;
+
+import com.amazonaws.AmazonServiceException;
+import com.amazonaws.SdkClientException;
+import com.amazonaws.services.lambda.runtime.events.SQSEvent;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.ValueSource;
+import software.amazon.awssdk.services.sqs.SqsClient;
+
+import static com.amazonaws.services.lambda.runtime.events.SQSEvent.SQSMessage;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
+
+class PowertoolsSqsBatchProcessorTest {
+
+    private final SqsClient sqsClient = mock(SqsClient.class);
+    private final SqsClient interactionClient = mock(SqsClient.class);
+    private static final ObjectMapper MAPPER = new ObjectMapper();
+    private SQSEvent event;
+
+    @BeforeEach
+    void setUp() throws IOException {
+        event = MAPPER.readValue(this.getClass().getResource("/sampleSqsBatchEvent.json"), SQSEvent.class);
+        PowertoolsSqs.defaultSqsClient(sqsClient);
+    }
+
+    @Test
+    void shouldBatchProcessAndNotDeleteMessagesWhenAllSuccess() {
+        List<String> returnValues = PowertoolsSqs.partialBatchProcessor(event, false, (message) -> {
+            interactionClient.listQueues();
+            return "Success";
+        });
+
+        assertThat(returnValues)
+                .hasSize(2)
+                .containsExactly("Success", "Success");
+
+        verify(interactionClient, times(2)).listQueues();
+        verifyNoInteractions(sqsClient);
+    }
+
+    @ParameterizedTest
+    @ValueSource(classes = {SampleInnerSqsHandler.class, SampleSqsHandler.class})
+    void shouldBatchProcessViaClassAndNotDeleteMessagesWhenAllSuccess(Class<? extends SqsMessageHandler<String>> handler) {
+        List<String> returnValues = PowertoolsSqs.partialBatchProcessor(event, false, handler);
+
+        assertThat(returnValues)
+                .hasSize(2)
+                .containsExactly("0", "1");
+    }
+
+    private static Stream<Arguments> exception() {
+        return Stream.of(Arguments.of(new AmazonServiceException("Service Exception")),
+                Arguments.of(new SdkClientException("Client Exception")));
+    }
+
+    public class SampleInnerSqsHandler implements SqsMessageHandler<String> {
+        private int counter;
+
+        @Override
+        public String process(SQSMessage message) {
+            return String.valueOf(counter++);
+        }
+    }
+}

powertools-sqs/src/test/java/software/amazon/lambda/powertools/sqs/PowertoolsSqsTest.java renamed to powertools-sqs/src/test/java/software/amazon/lambda/powertools/sqs/PowertoolsSqsLargeMessageTest.java

@@ -35,7 +35,7 @@
 import static org.mockito.Mockito.when;
 import static org.mockito.MockitoAnnotations.openMocks;
 
-class PowertoolsSqsTest {
+class PowertoolsSqsLargeMessageTest {
 
     @Mock
     private AmazonS3 amazonS3;

powertools-sqs/src/test/java/software/amazon/lambda/powertools/sqs/SampleSqsHandler.java

@@ -0,0 +1,12 @@
+package software.amazon.lambda.powertools.sqs;
+
+import com.amazonaws.services.lambda.runtime.events.SQSEvent;
+
+public class SampleSqsHandler implements SqsMessageHandler<String> {
+    private int counter;
+
+    @Override
+    public String process(SQSEvent.SQSMessage message) {
+        return String.valueOf(counter++);
+    }
+}

powertools-sqs/src/test/resources/sampleSqsBatchEvent.json

@@ -0,0 +1,36 @@
+{
+  "records": [
+    {
+      "messageId": "059f36b4-87a3-44ab-83d2-661975830a7d",
+      "receiptHandle": "AQEBwJnKyrHigUMZj6rYigCgxlaS3SLy0a...",
+      "body": "Test message.",
+      "attributes": {
+        "ApproximateReceiveCount": "1",
+        "SentTimestamp": "1545082649183",
+        "SenderId": "AIDAIENQZJOLO23YVJ4VO",
+        "ApproximateFirstReceiveTimestamp": "1545082649185"
+      },
+      "messageAttributes": {},
+      "md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b3",
+      "eventSource": "aws:sqs",
+      "eventSourceArn": "arn:aws:sqs:us-east-2:123456789012:my-queue",
+      "awsRegion": "us-east-2"
+    },
+    {
+      "messageId": "2e1424d4-f796-459a-8184-9c92662be6da",
+      "receiptHandle": "AQEBzWwaftRI0KuVm4tP+/7q1rGgNqicHq...",
+      "body": "Test message.",
+      "attributes": {
+        "ApproximateReceiveCount": "1",
+        "SentTimestamp": "1545082650636",
+        "SenderId": "AIDAIENQZJOLO23YVJ4VO",
+        "ApproximateFirstReceiveTimestamp": "1545082650649"
+      },
+      "messageAttributes": {},
+      "md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b3",
+      "eventSource": "aws:sqs",
+      "eventSourceArn": "arn:aws:sqs:us-east-2:123456789012:my-queue",
+      "awsRegion": "us-east-2"
+    }
+  ]
+}