25
25
import static com .github .tomakehurst .wiremock .client .WireMock .urlEqualTo ;
26
26
import static com .github .tomakehurst .wiremock .client .WireMock .verify ;
27
27
import static org .junit .jupiter .api .Assertions .assertArrayEquals ;
28
+ import static org .junit .jupiter .api .Assertions .assertEquals ;
29
+ import static org .junit .jupiter .api .Assertions .assertNotEquals ;
30
+ import static org .junit .jupiter .api .Assertions .assertNotNull ;
31
+ import static org .junit .jupiter .api .Assertions .assertThrows ;
28
32
29
33
import com .github .tomakehurst .wiremock .junit5 .WireMockRuntimeInfo ;
30
34
import com .github .tomakehurst .wiremock .junit5 .WireMockTest ;
36
40
import java .util .Random ;
37
41
import java .util .UUID ;
38
42
import java .util .concurrent .CompletableFuture ;
43
+ import java .util .concurrent .CompletionException ;
39
44
import org .junit .jupiter .api .BeforeEach ;
40
45
import org .junit .jupiter .api .Test ;
41
46
import software .amazon .awssdk .auth .credentials .AwsBasicCredentials ;
42
47
import software .amazon .awssdk .auth .credentials .StaticCredentialsProvider ;
43
48
import software .amazon .awssdk .awscore .retry .AwsRetryStrategy ;
44
49
import software .amazon .awssdk .core .ResponseBytes ;
45
50
import software .amazon .awssdk .core .async .AsyncResponseTransformer ;
51
+ import software .amazon .awssdk .core .interceptor .Context ;
52
+ import software .amazon .awssdk .core .interceptor .ExecutionAttributes ;
53
+ import software .amazon .awssdk .core .interceptor .ExecutionInterceptor ;
54
+ import software .amazon .awssdk .http .SdkHttpResponse ;
46
55
import software .amazon .awssdk .http .nio .netty .NettyNioAsyncHttpClient ;
47
56
import software .amazon .awssdk .regions .Region ;
48
57
import software .amazon .awssdk .services .s3 .S3AsyncClient ;
@@ -60,11 +69,14 @@ public class S3MultipartClientGetObjectWiremockTest {
60
69
+ "</Error>" ;
61
70
public static final String BUCKET = "Example-Bucket" ;
62
71
public static final String KEY = "Key" ;
72
+ private static final int MAX_ATTEMPTS = 7 ;
73
+ private static final CapturingInterceptor capturingInterceptor = new CapturingInterceptor ();
63
74
64
75
private S3AsyncClient multipartClient ;
65
76
66
77
@ BeforeEach
67
78
public void setup (WireMockRuntimeInfo wm ) {
79
+ capturingInterceptor .clear ();
68
80
multipartClient = S3AsyncClient .builder ()
69
81
.region (Region .US_EAST_1 )
70
82
.endpointOverride (URI .create (wm .getHttpBaseUrl ()))
@@ -73,9 +85,10 @@ public void setup(WireMockRuntimeInfo wm) {
73
85
.credentialsProvider (StaticCredentialsProvider .create (AwsBasicCredentials .create ("key" , "secret" )))
74
86
.overrideConfiguration (
75
87
o -> o .retryStrategy (AwsRetryStrategy .standardRetryStrategy ().toBuilder ()
76
- .maxAttempts (10 )
88
+ .maxAttempts (MAX_ATTEMPTS )
77
89
.circuitBreakerEnabled (false )
78
- .build ()))
90
+ .build ())
91
+ .addExecutionInterceptor (capturingInterceptor ))
79
92
.build ();
80
93
}
81
94
@@ -127,67 +140,47 @@ public void stub_503_then_200_multipleTimes() {
127
140
CompletableFuture .allOf (futures .toArray (new CompletableFuture [0 ])).join ();
128
141
}
129
142
130
- // TODO - assert first request ID not reused on retries
131
- /*@Test
143
+ @ Test
132
144
public void stub_503_only (WireMockRuntimeInfo wm ) {
145
+ String firstRequestId = UUID .randomUUID ().toString ();
146
+ String secondRequestId = UUID .randomUUID ().toString ();
147
+
133
148
stubFor (any (anyUrl ())
134
149
.inScenario ("errors" )
135
150
.whenScenarioStateIs (Scenario .STARTED )
136
151
.willReturn (aResponse ()
137
- .withHeader("x-amz-request-id", String.valueOf(UUID.randomUUID()) )
152
+ .withHeader ("x-amz-request-id" , firstRequestId )
138
153
.withStatus (503 )
139
154
.withBody (ERROR_BODY ))
140
155
.willSetStateTo ("SecondAttempt" ));
141
156
142
157
stubFor (any (anyUrl ())
143
158
.inScenario ("errors" )
144
159
.whenScenarioStateIs ("SecondAttempt" )
145
- .willReturn(aResponse().withHeader("x-amz-request-id", String.valueOf(UUID.randomUUID()))
146
- .withStatus(503)));
160
+ .willReturn (aResponse ()
161
+ .withHeader ("x-amz-request-id" , secondRequestId )
162
+ .withStatus (503 )));
147
163
148
- multipartClient.getObject(b -> b.bucket(BUCKET).key(KEY), AsyncResponseTransformer.toBytes()).join();
149
- }*/
164
+ assertThrows (CompletionException .class , () -> {
165
+ multipartClient .getObject (b -> b .bucket (BUCKET ).key (KEY ), AsyncResponseTransformer .toBytes ()).join ();
166
+ });
150
167
151
- private CompletableFuture < ResponseBytes < GetObjectResponse >> mock200Response ( S3AsyncClient s3Client , int runNumber ) {
152
- String runId = runNumber + " sucess" ;
168
+ List < SdkHttpResponse > responses = capturingInterceptor . getResponses ();
169
+ assertEquals ( MAX_ATTEMPTS , responses . size (), () -> String . format ( "Expected exactly %s responses" , MAX_ATTEMPTS )) ;
153
170
154
- stubFor (any (anyUrl ())
155
- .withHeader ("RunNum" , matching (runId ))
156
- .inScenario (runId )
157
- .whenScenarioStateIs (Scenario .STARTED )
158
- .willReturn (aResponse ().withStatus (200 )
159
- .withHeader ("x-amz-request-id" , String .valueOf (UUID .randomUUID ()))
160
- .withBody ("Hello World" )));
171
+ String actualFirstRequestId = responses .get (0 ).firstMatchingHeader ("x-amz-request-id" ).orElse (null );
172
+ String actualSecondRequestId = responses .get (1 ).firstMatchingHeader ("x-amz-request-id" ).orElse (null );
161
173
162
- return s3Client .getObject (r -> r .bucket (BUCKET ).key ("key" )
163
- .overrideConfiguration (c -> c .putHeader ("RunNum" , runId )),
164
- AsyncResponseTransformer .toBytes ());
165
- }
174
+ assertNotNull (actualFirstRequestId , "First response should have x-amz-request-id header" );
175
+ assertNotNull (actualSecondRequestId , "Second response should have x-amz-request-id header" );
166
176
167
- private CompletableFuture <ResponseBytes <GetObjectResponse >> mockRetryableErrorThen200Response (S3AsyncClient s3Client , int runNumber ) {
168
- String runId = String .valueOf (runNumber );
177
+ assertNotEquals (actualFirstRequestId , actualSecondRequestId , "First request ID should not be reused on retry" );
169
178
170
- stubFor (any (anyUrl ())
171
- .withHeader ("RunNum" , matching (runId ))
172
- .inScenario (runId )
173
- .whenScenarioStateIs (Scenario .STARTED )
174
- .willReturn (aResponse ()
175
- .withHeader ("x-amz-request-id" , String .valueOf (UUID .randomUUID ()))
176
- .withStatus (503 ).withBody (ERROR_BODY )
177
- )
178
- .willSetStateTo ("SecondAttempt" + runId ));
179
-
180
- stubFor (any (anyUrl ())
181
- .inScenario (runId )
182
- .withHeader ("RunNum" , matching (runId ))
183
- .whenScenarioStateIs ("SecondAttempt" + runId )
184
- .willReturn (aResponse ().withStatus (200 )
185
- .withHeader ("x-amz-request-id" , String .valueOf (UUID .randomUUID ()))
186
- .withBody ("Hello World" )));
179
+ assertEquals (firstRequestId , actualFirstRequestId , "First response should have expected request ID" );
180
+ assertEquals (secondRequestId , actualSecondRequestId , "Second response should have expected request ID" );
187
181
188
- return s3Client .getObject (r -> r .bucket (BUCKET ).key ("key" )
189
- .overrideConfiguration (c -> c .putHeader ("RunNum" , runId )),
190
- AsyncResponseTransformer .toBytes ());
182
+ assertEquals (503 , responses .get (0 ).statusCode ());
183
+ assertEquals (503 , responses .get (1 ).statusCode ());
191
184
}
192
185
193
186
@ Test
@@ -310,4 +303,63 @@ public void multipleParts_503OnFirstPart_then_200s() {
310
303
verify (1 , getRequestedFor (urlEqualTo (String .format ("/%s/%s?partNumber=2" , BUCKET , KEY ))));
311
304
verify (1 , getRequestedFor (urlEqualTo (String .format ("/%s/%s?partNumber=3" , BUCKET , KEY ))));
312
305
}
306
+
307
+ private CompletableFuture <ResponseBytes <GetObjectResponse >> mock200Response (S3AsyncClient s3Client , int runNumber ) {
308
+ String runId = runNumber + " sucess" ;
309
+
310
+ stubFor (any (anyUrl ())
311
+ .withHeader ("RunNum" , matching (runId ))
312
+ .inScenario (runId )
313
+ .whenScenarioStateIs (Scenario .STARTED )
314
+ .willReturn (aResponse ().withStatus (200 )
315
+ .withHeader ("x-amz-request-id" , String .valueOf (UUID .randomUUID ()))
316
+ .withBody ("Hello World" )));
317
+
318
+ return s3Client .getObject (r -> r .bucket (BUCKET ).key ("key" )
319
+ .overrideConfiguration (c -> c .putHeader ("RunNum" , runId )),
320
+ AsyncResponseTransformer .toBytes ());
321
+ }
322
+
323
+ private CompletableFuture <ResponseBytes <GetObjectResponse >> mockRetryableErrorThen200Response (S3AsyncClient s3Client , int runNumber ) {
324
+ String runId = String .valueOf (runNumber );
325
+
326
+ stubFor (any (anyUrl ())
327
+ .withHeader ("RunNum" , matching (runId ))
328
+ .inScenario (runId )
329
+ .whenScenarioStateIs (Scenario .STARTED )
330
+ .willReturn (aResponse ()
331
+ .withHeader ("x-amz-request-id" , String .valueOf (UUID .randomUUID ()))
332
+ .withStatus (503 ).withBody (ERROR_BODY )
333
+ )
334
+ .willSetStateTo ("SecondAttempt" + runId ));
335
+
336
+ stubFor (any (anyUrl ())
337
+ .inScenario (runId )
338
+ .withHeader ("RunNum" , matching (runId ))
339
+ .whenScenarioStateIs ("SecondAttempt" + runId )
340
+ .willReturn (aResponse ().withStatus (200 )
341
+ .withHeader ("x-amz-request-id" , String .valueOf (UUID .randomUUID ()))
342
+ .withBody ("Hello World" )));
343
+
344
+ return s3Client .getObject (r -> r .bucket (BUCKET ).key ("key" )
345
+ .overrideConfiguration (c -> c .putHeader ("RunNum" , runId )),
346
+ AsyncResponseTransformer .toBytes ());
347
+ }
348
+
349
+ static class CapturingInterceptor implements ExecutionInterceptor {
350
+ private final List <SdkHttpResponse > responses = new ArrayList <>();
351
+
352
+ @ Override
353
+ public void afterTransmission (Context .AfterTransmission context , ExecutionAttributes executionAttributes ) {
354
+ responses .add (context .httpResponse ());
355
+ }
356
+
357
+ public List <SdkHttpResponse > getResponses () {
358
+ return new ArrayList <>(responses );
359
+ }
360
+
361
+ public void clear () {
362
+ responses .clear ();
363
+ }
364
+ }
313
365
}
0 commit comments