Handle void response in injectCredentialsAndInvoke (#338)

* Overload function and use Consumer interface

Author: omkhegde

src/main/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxy.java

@@ -27,6 +27,7 @@
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.BiFunction;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
@@ -443,9 +444,24 @@ public ProgressEvent<ModelT, CallbackT> done(Function<ResponseT, ProgressEvent<M
         request.setRequestCredentialsProvider(v1CredentialsProvider);
 
         try {
-            ResultT respose = requestFunction.apply(request);
-            logRequestMetadata(request, respose);
-            return respose;
+            ResultT response = requestFunction.apply(request);
+            logRequestMetadata(request, response);
+            return response;
+        } catch (final Throwable e) {
+            loggerProxy.log(String.format("Failed to execute remote function: {%s}", e.getMessage()));
+            throw e;
+        } finally {
+            request.setRequestCredentialsProvider(null);
+        }
+    }
+
+    public <RequestT extends AmazonWebServiceRequest> void injectCredentialsAndInvoke(final RequestT request,
+                                                                                      final Consumer<RequestT> requestFunction) {
+
+        request.setRequestCredentialsProvider(v1CredentialsProvider);
+
+        try {
+            requestFunction.accept(request);
         } catch (final Throwable e) {
             loggerProxy.log(String.format("Failed to execute remote function: {%s}", e.getMessage()));
             throw e;

src/test/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxyTest.java

@@ -76,6 +76,12 @@ public class AmazonWebServicesClientProxyTest {
     //
     private final AwsSessionCredentials MockCreds = AwsSessionCredentials.create("accessKeyId", "secretKey", "token");
 
+    //
+    // Empty method for testing injectCredentialsAndInvoke with void returns
+    //
+    public static void dummyMethod(DescribeStackEventsRequest request) {
+    }
+
     @Test
     public void testInjectCredentialsAndInvoke() {
 
@@ -102,6 +108,23 @@ public void testInjectCredentialsAndInvoke() {
         assertThat(result).isEqualTo(expectedResult);
     }
 
+    @Test
+    public void testInjectCredentialsAndInvokeWithVoidFunction() {
+
+        final LoggerProxy loggerProxy = mock(LoggerProxy.class);
+        final Credentials credentials = new Credentials("accessKeyId", "secretAccessKey", "sessionToken");
+
+        final AmazonWebServicesClientProxy proxy = new AmazonWebServicesClientProxy(loggerProxy, credentials, () -> 1000L);
+
+        final DescribeStackEventsRequest request = mock(DescribeStackEventsRequest.class);
+
+        proxy.injectCredentialsAndInvoke(request, AmazonWebServicesClientProxyTest::dummyMethod);
+
+        // ensure credentials are injected and then removed
+        verify(request).setRequestCredentialsProvider(any(AWSStaticCredentialsProvider.class));
+        verify(request).setRequestCredentialsProvider(eq(null));
+    }
+
     @Test
     public void testInjectCredentialsAndInvokeWithError() {