2525import static com .github .tomakehurst .wiremock .client .WireMock .urlEqualTo ;
2626import static com .github .tomakehurst .wiremock .client .WireMock .verify ;
2727import 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 ;
2832
2933import com .github .tomakehurst .wiremock .junit5 .WireMockRuntimeInfo ;
3034import com .github .tomakehurst .wiremock .junit5 .WireMockTest ;
3640import java .util .Random ;
3741import java .util .UUID ;
3842import java .util .concurrent .CompletableFuture ;
43+ import java .util .concurrent .CompletionException ;
3944import org .junit .jupiter .api .BeforeEach ;
4045import org .junit .jupiter .api .Test ;
4146import software .amazon .awssdk .auth .credentials .AwsBasicCredentials ;
4247import software .amazon .awssdk .auth .credentials .StaticCredentialsProvider ;
4348import software .amazon .awssdk .awscore .retry .AwsRetryStrategy ;
4449import software .amazon .awssdk .core .ResponseBytes ;
4550import 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 ;
4655import software .amazon .awssdk .http .nio .netty .NettyNioAsyncHttpClient ;
4756import software .amazon .awssdk .regions .Region ;
4857import software .amazon .awssdk .services .s3 .S3AsyncClient ;
@@ -60,11 +69,14 @@ public class S3MultipartClientGetObjectWiremockTest {
6069 + "</Error>" ;
6170 public static final String BUCKET = "Example-Bucket" ;
6271 public static final String KEY = "Key" ;
72+ private static final int MAX_ATTEMPTS = 7 ;
73+ private static final CapturingInterceptor capturingInterceptor = new CapturingInterceptor ();
6374
6475 private S3AsyncClient multipartClient ;
6576
6677 @ BeforeEach
6778 public void setup (WireMockRuntimeInfo wm ) {
79+ capturingInterceptor .clear ();
6880 multipartClient = S3AsyncClient .builder ()
6981 .region (Region .US_EAST_1 )
7082 .endpointOverride (URI .create (wm .getHttpBaseUrl ()))
@@ -73,9 +85,10 @@ public void setup(WireMockRuntimeInfo wm) {
7385 .credentialsProvider (StaticCredentialsProvider .create (AwsBasicCredentials .create ("key" , "secret" )))
7486 .overrideConfiguration (
7587 o -> o .retryStrategy (AwsRetryStrategy .standardRetryStrategy ().toBuilder ()
76- .maxAttempts (10 )
88+ .maxAttempts (MAX_ATTEMPTS )
7789 .circuitBreakerEnabled (false )
78- .build ()))
90+ .build ())
91+ .addExecutionInterceptor (capturingInterceptor ))
7992 .build ();
8093 }
8194
@@ -127,67 +140,47 @@ public void stub_503_then_200_multipleTimes() {
127140 CompletableFuture .allOf (futures .toArray (new CompletableFuture [0 ])).join ();
128141 }
129142
130- // TODO - assert first request ID not reused on retries
131- /*@Test
143+ @ Test
132144 public void stub_503_only (WireMockRuntimeInfo wm ) {
145+ String firstRequestId = UUID .randomUUID ().toString ();
146+ String secondRequestId = UUID .randomUUID ().toString ();
147+
133148 stubFor (any (anyUrl ())
134149 .inScenario ("errors" )
135150 .whenScenarioStateIs (Scenario .STARTED )
136151 .willReturn (aResponse ()
137- .withHeader("x-amz-request-id", String.valueOf(UUID.randomUUID()) )
152+ .withHeader ("x-amz-request-id" , firstRequestId )
138153 .withStatus (503 )
139154 .withBody (ERROR_BODY ))
140155 .willSetStateTo ("SecondAttempt" ));
141156
142157 stubFor (any (anyUrl ())
143158 .inScenario ("errors" )
144159 .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 )));
147163
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+ });
150167
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 )) ;
153170
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 );
161173
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" );
166176
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" );
169178
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" );
187181
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 ());
191184 }
192185
193186 @ Test
@@ -310,4 +303,63 @@ public void multipleParts_503OnFirstPart_then_200s() {
310303 verify (1 , getRequestedFor (urlEqualTo (String .format ("/%s/%s?partNumber=2" , BUCKET , KEY ))));
311304 verify (1 , getRequestedFor (urlEqualTo (String .format ("/%s/%s?partNumber=3" , BUCKET , KEY ))));
312305 }
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+ }
313365}
0 commit comments