diff --git a/src/main/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxy.java b/src/main/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxy.java index f2683809..420b1cbf 100644 --- a/src/main/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxy.java +++ b/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 done(Function void injectCredentialsAndInvoke(final RequestT request, + final Consumer 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; diff --git a/src/test/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxyTest.java b/src/test/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxyTest.java index 9cbf0c32..d60b7d4d 100644 --- a/src/test/java/software/amazon/cloudformation/proxy/AmazonWebServicesClientProxyTest.java +++ b/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() {