Skip to content

Commit a31e120

Browse files
xiwhuangxiwhuang
authored
Add proxy methods to support invocation for inputstream/bytes (#259)
* Add proxy methods to support invocation for inputstream/bytes * Add two proxy methods in code-gen template
1 parent a64e61e commit a31e120

File tree

5 files changed

+266
-0
lines changed

5 files changed

+266
-0
lines changed

‎pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,11 @@
155155
<groupId>software.amazon.awssdk</groupId>
156156
<artifactId>apache-client</artifactId>
157157
</dependency>
158+
<dependency>
159+
<groupId>software.amazon.awssdk</groupId>
160+
<artifactId>s3</artifactId>
161+
<scope>test</scope>
162+
</dependency>
158163

159164
<!-- https://mvnrepository.com/artifact/commons-io/commons-io -->
160165
<dependency>

‎python/rpdk/java/templates/init/guided_aws/AbstractTestBase.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
import software.amazon.awssdk.awscore.AwsRequest;
66
import software.amazon.awssdk.awscore.AwsResponse;
77
import software.amazon.awssdk.core.SdkClient;
8+
import software.amazon.awssdk.core.ResponseBytes;
9+
import software.amazon.awssdk.core.ResponseInputStream;
810
import software.amazon.awssdk.core.pagination.sync.SdkIterable;
911
import software.amazon.cloudformation.proxy.AmazonWebServicesClientProxy;
1012
import software.amazon.cloudformation.proxy.Credentials;
@@ -43,6 +45,18 @@ static ProxyClient<SdkClient> MOCK_PROXY(
4345
return proxy.injectCredentialsAndInvokeIterableV2(request, requestFunction);
4446
}
4547

48+
@Override
49+
public <RequestT extends AwsRequest, ResponseT extends AwsResponse> ResponseInputStream<ResponseT>
50+
injectCredentialsAndInvokeV2InputStream(RequestT requestT, Function<RequestT, ResponseInputStream<ResponseT>> function) {
51+
throw new UnsupportedOperationException();
52+
}
53+
54+
@Override
55+
public <RequestT extends AwsRequest, ResponseT extends AwsResponse> ResponseBytes<ResponseT>
56+
injectCredentialsAndInvokeV2Bytes(RequestT requestT, Function<RequestT, ResponseBytes<ResponseT>> function) {
57+
throw new UnsupportedOperationException();
58+
}
59+
4660
@Override
4761
public SdkClient client() {
4862
return sdkClient;

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

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@
4040
import software.amazon.awssdk.awscore.AwsResponse;
4141
import software.amazon.awssdk.awscore.exception.AwsErrorDetails;
4242
import software.amazon.awssdk.awscore.exception.AwsServiceException;
43+
import software.amazon.awssdk.core.ResponseBytes;
44+
import software.amazon.awssdk.core.ResponseInputStream;
4345
import software.amazon.awssdk.core.SdkClient;
4446
import software.amazon.awssdk.core.exception.NonRetryableException;
4547
import software.amazon.awssdk.core.exception.RetryableException;
@@ -124,6 +126,22 @@ public <ClientT> ProxyClient<ClientT> newProxy(@Nonnull Supplier<ClientT> client
124126
return AmazonWebServicesClientProxy.this.injectCredentialsAndInvokeIterableV2(request, requestFunction);
125127
}
126128

129+
@Override
130+
public <RequestT extends AwsRequest, ResponseT extends AwsResponse>
131+
ResponseInputStream<ResponseT>
132+
injectCredentialsAndInvokeV2InputStream(RequestT request,
133+
Function<RequestT, ResponseInputStream<ResponseT>> requestFunction) {
134+
return AmazonWebServicesClientProxy.this.injectCredentialsAndInvokeV2InputStream(request, requestFunction);
135+
}
136+
137+
@Override
138+
public <RequestT extends AwsRequest, ResponseT extends AwsResponse>
139+
ResponseBytes<ResponseT>
140+
injectCredentialsAndInvokeV2Bytes(RequestT request,
141+
Function<RequestT, ResponseBytes<ResponseT>> requestFunction) {
142+
return AmazonWebServicesClientProxy.this.injectCredentialsAndInvokeV2Bytes(request, requestFunction);
143+
}
144+
127145
@Override
128146
public ClientT client() {
129147
return client.get();
@@ -512,6 +530,44 @@ public final long getRemainingTimeInMillis() {
512530
}
513531
}
514532

533+
public <RequestT extends AwsRequest, ResultT extends AwsResponse>
534+
ResponseInputStream<ResultT>
535+
injectCredentialsAndInvokeV2InputStream(final RequestT request,
536+
final Function<RequestT, ResponseInputStream<ResultT>> requestFunction) {
537+
538+
AwsRequestOverrideConfiguration overrideConfiguration = AwsRequestOverrideConfiguration.builder()
539+
.credentialsProvider(v2CredentialsProvider).build();
540+
541+
@SuppressWarnings("unchecked")
542+
RequestT wrappedRequest = (RequestT) request.toBuilder().overrideConfiguration(overrideConfiguration).build();
543+
544+
try {
545+
return requestFunction.apply(wrappedRequest);
546+
} catch (final Throwable e) {
547+
loggerProxy.log(String.format("Failed to execute remote function: {%s}", e.getMessage()));
548+
throw e;
549+
}
550+
}
551+
552+
public <RequestT extends AwsRequest, ResultT extends AwsResponse>
553+
ResponseBytes<ResultT>
554+
injectCredentialsAndInvokeV2Bytes(final RequestT request,
555+
final Function<RequestT, ResponseBytes<ResultT>> requestFunction) {
556+
557+
AwsRequestOverrideConfiguration overrideConfiguration = AwsRequestOverrideConfiguration.builder()
558+
.credentialsProvider(v2CredentialsProvider).build();
559+
560+
@SuppressWarnings("unchecked")
561+
RequestT wrappedRequest = (RequestT) request.toBuilder().overrideConfiguration(overrideConfiguration).build();
562+
563+
try {
564+
return requestFunction.apply(wrappedRequest);
565+
} catch (final Throwable e) {
566+
loggerProxy.log(String.format("Failed to execute remote function: {%s}", e.getMessage()));
567+
throw e;
568+
}
569+
}
570+
515571
public <RequestT, ClientT, ModelT, CallbackT extends StdCallbackContext>
516572
ProgressEvent<ModelT, CallbackT>
517573
defaultHandler(RequestT request, Exception e, ClientT client, ModelT model, CallbackT context) throws Exception {

‎src/main/java/software/amazon/cloudformation/proxy/ProxyClient.java‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
import java.util.function.Function;
1919
import software.amazon.awssdk.awscore.AwsRequest;
2020
import software.amazon.awssdk.awscore.AwsResponse;
21+
import software.amazon.awssdk.core.ResponseBytes;
22+
import software.amazon.awssdk.core.ResponseInputStream;
2123
import software.amazon.awssdk.core.pagination.sync.SdkIterable;
2224

2325
/**
@@ -94,6 +96,45 @@ public interface ProxyClient<ClientT> {
9496
IterableT
9597
injectCredentialsAndInvokeIterableV2(RequestT request, Function<RequestT, IterableT> requestFunction);
9698

99+
/**
100+
* This is a synchronous version of making API calls which implement
101+
* ResponseInputStream in the SDKv2
102+
*
103+
* @param request, the AWS service request that we need to make
104+
* @param requestFunction, this is a Lambda closure that provide the actual API
105+
* that needs to be invoked.
106+
* @param <RequestT> the request type
107+
* @param <ResponseT> the response from the request
108+
* @return the response if successful. Else it will propagate all
109+
* {@link software.amazon.awssdk.awscore.exception.AwsServiceException}
110+
* that is thrown or
111+
* {@link software.amazon.awssdk.core.exception.SdkClientException} if
112+
* there is client side problem
113+
*/
114+
<RequestT extends AwsRequest, ResponseT extends AwsResponse>
115+
ResponseInputStream<ResponseT>
116+
injectCredentialsAndInvokeV2InputStream(RequestT request,
117+
Function<RequestT, ResponseInputStream<ResponseT>> requestFunction);
118+
119+
/**
120+
* This is a synchronous version of making API calls which implement
121+
* ResponseBytes in the SDKv2
122+
*
123+
* @param request, the AWS service request that we need to make
124+
* @param requestFunction, this is a Lambda closure that provide the actual API
125+
* that needs to be invoked.
126+
* @param <RequestT> the request type
127+
* @param <ResponseT> the response from the request
128+
* @return the response if successful. Else it will propagate all
129+
* {@link software.amazon.awssdk.awscore.exception.AwsServiceException}
130+
* that is thrown or
131+
* {@link software.amazon.awssdk.core.exception.SdkClientException} if
132+
* there is client side problem
133+
*/
134+
<RequestT extends AwsRequest, ResponseT extends AwsResponse>
135+
ResponseBytes<ResponseT>
136+
injectCredentialsAndInvokeV2Bytes(RequestT request, Function<RequestT, ResponseBytes<ResponseT>> requestFunction);
137+
97138
/**
98139
* @return the actual AWS service client that we need to use to provide the
99140
* actual method we are going to call.

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

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020
import static org.junit.jupiter.api.Assertions.fail;
2121
import static org.mockito.ArgumentMatchers.any;
2222
import static org.mockito.ArgumentMatchers.eq;
23+
import static org.mockito.Mockito.doReturn;
24+
import static org.mockito.Mockito.doThrow;
2325
import static org.mockito.Mockito.mock;
2426
import static org.mockito.Mockito.times;
2527
import static org.mockito.Mockito.verify;
@@ -42,11 +44,15 @@
4244
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
4345
import software.amazon.awssdk.awscore.exception.AwsErrorDetails;
4446
import software.amazon.awssdk.awscore.exception.AwsServiceException;
47+
import software.amazon.awssdk.core.ResponseBytes;
48+
import software.amazon.awssdk.core.ResponseInputStream;
4549
import software.amazon.awssdk.core.exception.NonRetryableException;
4650
import software.amazon.awssdk.http.SdkHttpResponse;
4751
import software.amazon.awssdk.services.cloudformation.CloudFormationAsyncClient;
4852
import software.amazon.awssdk.services.cloudformation.CloudFormationClient;
4953
import software.amazon.awssdk.services.cloudformation.model.DescribeStackEventsResponse;
54+
import software.amazon.awssdk.services.s3.S3Client;
55+
import software.amazon.awssdk.services.s3.model.GetObjectResponse;
5056
import software.amazon.cloudformation.exceptions.ResourceAlreadyExistsException;
5157
import software.amazon.cloudformation.exceptions.TerminalException;
5258
import software.amazon.cloudformation.proxy.delay.Constant;
@@ -229,6 +235,150 @@ public void testInjectCredentialsAndInvokeV2Async_WithException() {
229235

230236
}
231237

238+
@Test
239+
public void testInjectCredentialsAndInvokeV2InputStream() {
240+
241+
final LoggerProxy loggerProxy = mock(LoggerProxy.class);
242+
final Credentials credentials = new Credentials("accessKeyId", "secretAccessKey", "sessionToken");
243+
final ResponseInputStream<?> responseInputStream = mock(ResponseInputStream.class);
244+
245+
final AmazonWebServicesClientProxy proxy = new AmazonWebServicesClientProxy(loggerProxy, credentials, () -> 1000L);
246+
247+
final software.amazon.awssdk.services.s3.model.GetObjectRequest wrappedRequest = mock(
248+
software.amazon.awssdk.services.s3.model.GetObjectRequest.class);
249+
250+
final software.amazon.awssdk.services.s3.model.GetObjectRequest.Builder builder = mock(
251+
software.amazon.awssdk.services.s3.model.GetObjectRequest.Builder.class);
252+
when(builder.overrideConfiguration(any(AwsRequestOverrideConfiguration.class))).thenReturn(builder);
253+
when(builder.build()).thenReturn(wrappedRequest);
254+
final software.amazon.awssdk.services.s3.model.GetObjectRequest request = mock(
255+
software.amazon.awssdk.services.s3.model.GetObjectRequest.class);
256+
when(request.toBuilder()).thenReturn(builder);
257+
258+
final S3Client client = mock(S3Client.class);
259+
260+
doReturn(responseInputStream).when(client)
261+
.getObject(any(software.amazon.awssdk.services.s3.model.GetObjectRequest.class));
262+
263+
final ResponseInputStream<
264+
GetObjectResponse> result = proxy.injectCredentialsAndInvokeV2InputStream(request, client::getObject);
265+
266+
// verify request is rebuilt for injection
267+
verify(request).toBuilder();
268+
269+
// verify the wrapped request is sent over the initiate
270+
verify(client).getObject(wrappedRequest);
271+
272+
// ensure the return type matches
273+
assertThat(result).isEqualTo(responseInputStream);
274+
}
275+
276+
@Test
277+
public void testInjectCredentialsAndInvokeV2InputStream_Exception() {
278+
279+
final LoggerProxy loggerProxy = mock(LoggerProxy.class);
280+
final Credentials credentials = new Credentials("accessKeyId", "secretAccessKey", "sessionToken");
281+
282+
final AmazonWebServicesClientProxy proxy = new AmazonWebServicesClientProxy(loggerProxy, credentials, () -> 1000L);
283+
284+
final software.amazon.awssdk.services.s3.model.GetObjectRequest wrappedRequest = mock(
285+
software.amazon.awssdk.services.s3.model.GetObjectRequest.class);
286+
287+
final software.amazon.awssdk.services.s3.model.GetObjectRequest.Builder builder = mock(
288+
software.amazon.awssdk.services.s3.model.GetObjectRequest.Builder.class);
289+
when(builder.overrideConfiguration(any(AwsRequestOverrideConfiguration.class))).thenReturn(builder);
290+
when(builder.build()).thenReturn(wrappedRequest);
291+
final software.amazon.awssdk.services.s3.model.GetObjectRequest request = mock(
292+
software.amazon.awssdk.services.s3.model.GetObjectRequest.class);
293+
when(request.toBuilder()).thenReturn(builder);
294+
295+
final S3Client client = mock(S3Client.class);
296+
297+
doThrow(new TerminalException(new RuntimeException("Sorry"))).when(client)
298+
.getObject(any(software.amazon.awssdk.services.s3.model.GetObjectRequest.class));
299+
300+
assertThrows(RuntimeException.class, () -> proxy.injectCredentialsAndInvokeV2InputStream(request, client::getObject),
301+
"Expected Runtime Exception.");
302+
303+
// verify request is rebuilt for injection
304+
verify(request).toBuilder();
305+
306+
// verify the wrapped request is sent over the initiate
307+
verify(client).getObject(wrappedRequest);
308+
}
309+
310+
@Test
311+
public void testInjectCredentialsAndInvokeV2Bytes() {
312+
313+
final LoggerProxy loggerProxy = mock(LoggerProxy.class);
314+
final Credentials credentials = new Credentials("accessKeyId", "secretAccessKey", "sessionToken");
315+
final ResponseBytes<?> responseBytes = mock(ResponseBytes.class);
316+
317+
final AmazonWebServicesClientProxy proxy = new AmazonWebServicesClientProxy(loggerProxy, credentials, () -> 1000L);
318+
319+
final software.amazon.awssdk.services.s3.model.GetObjectRequest wrappedRequest = mock(
320+
software.amazon.awssdk.services.s3.model.GetObjectRequest.class);
321+
322+
final software.amazon.awssdk.services.s3.model.GetObjectRequest.Builder builder = mock(
323+
software.amazon.awssdk.services.s3.model.GetObjectRequest.Builder.class);
324+
when(builder.overrideConfiguration(any(AwsRequestOverrideConfiguration.class))).thenReturn(builder);
325+
when(builder.build()).thenReturn(wrappedRequest);
326+
final software.amazon.awssdk.services.s3.model.GetObjectRequest request = mock(
327+
software.amazon.awssdk.services.s3.model.GetObjectRequest.class);
328+
when(request.toBuilder()).thenReturn(builder);
329+
330+
final S3Client client = mock(S3Client.class);
331+
332+
doReturn(responseBytes).when(client)
333+
.getObjectAsBytes(any(software.amazon.awssdk.services.s3.model.GetObjectRequest.class));
334+
335+
final ResponseBytes<
336+
GetObjectResponse> result = proxy.injectCredentialsAndInvokeV2Bytes(request, client::getObjectAsBytes);
337+
338+
// verify request is rebuilt for injection
339+
verify(request).toBuilder();
340+
341+
// verify the wrapped request is sent over the initiate
342+
verify(client).getObjectAsBytes(wrappedRequest);
343+
344+
// ensure the return type matches
345+
assertThat(result).isEqualTo(responseBytes);
346+
}
347+
348+
@Test
349+
public void testInjectCredentialsAndInvokeV2Bytes_Exception() {
350+
351+
final LoggerProxy loggerProxy = mock(LoggerProxy.class);
352+
final Credentials credentials = new Credentials("accessKeyId", "secretAccessKey", "sessionToken");
353+
354+
final AmazonWebServicesClientProxy proxy = new AmazonWebServicesClientProxy(loggerProxy, credentials, () -> 1000L);
355+
356+
final software.amazon.awssdk.services.s3.model.GetObjectRequest wrappedRequest = mock(
357+
software.amazon.awssdk.services.s3.model.GetObjectRequest.class);
358+
359+
final software.amazon.awssdk.services.s3.model.GetObjectRequest.Builder builder = mock(
360+
software.amazon.awssdk.services.s3.model.GetObjectRequest.Builder.class);
361+
when(builder.overrideConfiguration(any(AwsRequestOverrideConfiguration.class))).thenReturn(builder);
362+
when(builder.build()).thenReturn(wrappedRequest);
363+
final software.amazon.awssdk.services.s3.model.GetObjectRequest request = mock(
364+
software.amazon.awssdk.services.s3.model.GetObjectRequest.class);
365+
when(request.toBuilder()).thenReturn(builder);
366+
367+
final S3Client client = mock(S3Client.class);
368+
369+
doThrow(new TerminalException(new RuntimeException("Sorry"))).when(client)
370+
.getObjectAsBytes(any(software.amazon.awssdk.services.s3.model.GetObjectRequest.class));
371+
372+
assertThrows(RuntimeException.class, () -> proxy.injectCredentialsAndInvokeV2Bytes(request, client::getObjectAsBytes),
373+
"Expected Runtime Exception.");
374+
375+
// verify request is rebuilt for injection
376+
verify(request).toBuilder();
377+
378+
// verify the wrapped request is sent over the initiate
379+
verify(client).getObjectAsBytes(wrappedRequest);
380+
}
381+
232382
private final Credentials MOCK = new Credentials("accessKeyId", "secretKey", "token");
233383

234384
@Test

0 commit comments

Comments
 (0)