Skip to content

Commit 54457df

Browse files
committed
Update mock error test
1 parent 4cdded1 commit 54457df

File tree

1 file changed

+96
-44
lines changed

1 file changed

+96
-44
lines changed

services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/S3MultipartClientGetObjectWiremockTest.java

Lines changed: 96 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,10 @@
2525
import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
2626
import static com.github.tomakehurst.wiremock.client.WireMock.verify;
2727
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;
2832

2933
import com.github.tomakehurst.wiremock.junit5.WireMockRuntimeInfo;
3034
import com.github.tomakehurst.wiremock.junit5.WireMockTest;
@@ -36,13 +40,18 @@
3640
import java.util.Random;
3741
import java.util.UUID;
3842
import java.util.concurrent.CompletableFuture;
43+
import java.util.concurrent.CompletionException;
3944
import org.junit.jupiter.api.BeforeEach;
4045
import org.junit.jupiter.api.Test;
4146
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
4247
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
4348
import software.amazon.awssdk.awscore.retry.AwsRetryStrategy;
4449
import software.amazon.awssdk.core.ResponseBytes;
4550
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;
4655
import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient;
4756
import software.amazon.awssdk.regions.Region;
4857
import 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

Comments
 (0)