Integration test added

Author: Fred1155

services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/MultipartDownloaderSubscriber.java

@@ -138,7 +138,7 @@ public void onNext(AsyncResponseTransformer<GetObjectResponse, GetObjectResponse
         getObjectFuture.whenComplete((response, error) -> {
             if (error != null) {
                 log.debug(() -> "Error encountered during GetObjectRequest with partNumber=" + nextPartToGet);
-                onError(error);
+                handleError(error);
                 return;
             }
             requestMoreIfNeeded(response);
@@ -174,12 +174,20 @@ private void requestMoreIfNeeded(GetObjectResponse response) {
             if (totalParts != null && totalParts > 1 && totalComplete < totalParts) {
                 subscription.request(1);
             } else {
+                validatePartsCount();
                 log.debug(() -> String.format("Completing multipart download after a total of %d parts downloaded.", totalParts));
                 subscription.cancel();
             }
         }
     }
 
+    /**
+     * The method used by the Subscriber itself when error occured.
+     */
+    private void handleError(Throwable t) {
+        onError(t);
+    }
+
     @Override
     public void onError(Throwable t) {
         CompletableFuture<GetObjectResponse> partFuture;
@@ -191,7 +199,6 @@ public void onError(Throwable t) {
 
     @Override
     public void onComplete() {
-        validatePartsCount();
         future.complete(null);
     }
 
@@ -213,9 +220,8 @@ private void validatePartsCount() {
         if (totalParts != null && actualGetCount != totalParts) {
             String errorMessage = String.format("PartsCount validation failed. Expected %d, downloaded %d parts.", totalParts,
                                                 actualGetCount);
-            subscription.cancel();
             SdkClientException exception = SdkClientException.create(errorMessage);
-            onError(exception);
+            handleError(exception);
         }
     }
 }

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

@@ -34,6 +34,7 @@
 import static org.junit.jupiter.params.provider.Arguments.arguments;
 import static software.amazon.awssdk.services.s3.internal.multipart.utils.MultipartDownloadTestUtils.internalErrorBody;
 import static software.amazon.awssdk.services.s3.internal.multipart.utils.MultipartDownloadTestUtils.transformersSuppliers;
+import static software.amazon.awssdk.services.s3.multipart.S3MultipartExecutionAttribute.MULTIPART_DOWNLOAD_RESUME_CONTEXT;
 
 import com.github.tomakehurst.wiremock.http.Fault;
 import com.github.tomakehurst.wiremock.junit5.WireMockRuntimeInfo;
@@ -58,6 +59,7 @@
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
 import software.amazon.awssdk.core.SplittingTransformerConfiguration;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
+import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.core.internal.async.ByteArrayAsyncResponseTransformer;
 import software.amazon.awssdk.core.internal.async.FileAsyncResponseTransformer;
 import software.amazon.awssdk.core.internal.async.InputStreamResponseTransformer;
@@ -67,6 +69,7 @@
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.internal.multipart.utils.MultipartDownloadTestUtils;
+import software.amazon.awssdk.services.s3.model.GetObjectRequest;
 import software.amazon.awssdk.services.s3.model.GetObjectResponse;
 import software.amazon.awssdk.services.s3.model.S3Exception;
 import software.amazon.awssdk.services.s3.utils.AsyncResponseTransformerTestSupplier;
@@ -144,6 +147,47 @@ public <T> void errorOnThirdPart_shouldCompleteExceptionallyOnlyPartsGreaterThan
         }
     }
 
+    @ParameterizedTest
+    @MethodSource("partSizeAndTransformerParams")
+    public <T> void partCountValidationFailure_shouldThrowException(
+        AsyncResponseTransformerTestSupplier<T> supplier,
+        int partSize) {
+
+        // To trigger the partCount failure, the resumeContext is used to initialize the actualGetCount larger than the
+        // totalPart number set in the response. This won't happen in real scenario, just to test if the error can be surfaced
+        // to the user if the validation fails.
+        MultipartDownloadResumeContext resumeContext = new MultipartDownloadResumeContext();
+        resumeContext.addCompletedPart(1);
+        resumeContext.addCompletedPart(2);
+        resumeContext.addCompletedPart(3);
+        resumeContext.addToBytesToLastCompletedParts(3 * partSize);
+
+        GetObjectRequest request = GetObjectRequest.builder()
+                                                   .bucket(BUCKET)
+                                                   .key(KEY)
+                                                   .overrideConfiguration(config -> config
+                                                       .putExecutionAttribute(
+                                                           MULTIPART_DOWNLOAD_RESUME_CONTEXT,
+                                                           resumeContext))
+                                                   .build();
+
+        util.stubForPart(BUCKET, KEY, 4, 2, partSize);
+
+        // Skip the lazy transformer since the error won't surface unless the content is consumed
+        AsyncResponseTransformer<GetObjectResponse, T> transformer = supplier.transformer();
+        if (transformer instanceof InputStreamResponseTransformer || transformer instanceof PublisherAsyncResponseTransformer) {
+            return;
+        }
+
+        assertThatThrownBy(() -> {
+            T res = multipartClient.getObject(request, transformer).join();
+            supplier.body(res);
+        }).isInstanceOf(CompletionException.class)
+          .hasCauseInstanceOf(SdkClientException.class)
+          .hasMessageContaining("PartsCount validation failed. Expected 2, downloaded 4 parts");
+
+    }
+
     @ParameterizedTest
     @MethodSource("nonRetryableResponseTransformers")
     public <T> void errorOnFirstPart_shouldFail(AsyncResponseTransformerTestSupplier<T> supplier) {