capture client wip

Author: ibrandes

sdk/storage/azure-storage-blob/src/test/java/com/azure/storage/blob/contentValidation/BlobMessageEncoderUploadTests.java

@@ -3,17 +3,21 @@
 
 package com.azure.storage.blob.contentValidation;
 
+import com.azure.core.http.HttpMethod;
+import com.azure.core.http.HttpRequest;
 import com.azure.core.http.rest.Response;
 import com.azure.core.util.BinaryData;
 import com.azure.core.util.Context;
 import com.azure.core.util.ProgressListener;
 import com.azure.storage.blob.BlobClient;
+import com.azure.storage.blob.BlobClientBuilder;
 import com.azure.storage.blob.BlobTestBase;
 import com.azure.storage.blob.models.BlockBlobItem;
 import com.azure.storage.blob.models.ParallelTransferOptions;
 import com.azure.storage.blob.options.BlobParallelUploadOptions;
 import com.azure.storage.common.implementation.Constants;
 import com.azure.storage.common.implementation.structuredmessage.StorageChecksumAlgorithm;
+import com.azure.storage.common.test.shared.StorageCommonTestUtils;
 import com.azure.storage.common.test.shared.extensions.LiveOnly;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Disabled;
@@ -23,6 +27,8 @@
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.stream.Collectors;
 
 import static com.azure.storage.common.implementation.Constants.HeaderConstants.CONTENT_CRC64_HEADER_NAME;
 import static com.azure.storage.common.implementation.Constants.HeaderConstants.STRUCTURED_BODY_TYPE_HEADER_NAME;
@@ -32,6 +38,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 @LiveOnly
 public class BlobMessageEncoderUploadTests extends BlobTestBase {
@@ -88,6 +95,40 @@ public void uploadBinaryDataChunkedStructMess() {
         assertArrayEquals(data, outStream.toByteArray());
     }
 
+    @Test
+    void verifyUploadUsingChunkedStructMess() {
+        byte[] data = getRandomByteArray(Constants.MB * 10);
+
+        BlobParallelUploadOptions options = new BlobParallelUploadOptions(BinaryData.fromBytes(data))
+            .setStorageChecksumAlgorithm(StorageChecksumAlgorithm.AUTO)
+            .setParallelTransferOptions(
+                new ParallelTransferOptions().setMaxSingleUploadSizeLong((long) Constants.MB * 2)
+                    .setBlockSizeLong((long) Constants.MB * 2));
+
+        StorageCommonTestUtils.CaptureAllRequestsHttpClient httpClient
+            = new StorageCommonTestUtils.CaptureAllRequestsHttpClient(
+            StorageCommonTestUtils.getHttpClient(interceptorManager));
+        BlobClient capturingBlobClient = new BlobClientBuilder().endpoint(bc.getBlobUrl())
+            .credential(ENVIRONMENT.getPrimaryAccount().getCredential())
+            .httpClient(httpClient)
+            .buildClient();
+
+        capturingBlobClient.uploadWithResponse(options, null, Context.NONE);
+
+        List<HttpRequest> putRequests = httpClient.getCapturedRequests()
+            .stream()
+            .filter(r -> r.getHttpMethod() == HttpMethod.PUT)
+            .collect(Collectors.toList());
+
+        // Multipart upload should have at least 2 PUTs (staged + final)
+        assertTrue(putRequests.size() >= 2, "Expected multipart upload with multiple PUT requests.");
+
+        for (int i = 0; i < putRequests.size() - 1; i++) {
+            String headerValue = putRequests.get(i).getHeaders().getValue(STRUCTURED_BODY_TYPE_HEADER_NAME);
+            assertNotNull(headerValue, "Missing structured body header on staged PUT index " + i);
+        }
+    }
+
     @Test
     @Disabled
     public void uploadBinaryDataChunkedStructMessProgressListenerBehavior() {

sdk/storage/azure-storage-common/src/test-shared/java/com/azure/storage/common/test/shared/StorageCommonTestUtils.java

@@ -5,6 +5,8 @@
 import com.azure.core.client.traits.HttpTrait;
 import com.azure.core.credential.TokenCredential;
 import com.azure.core.http.HttpClient;
+import com.azure.core.http.HttpRequest;
+import com.azure.core.http.HttpResponse;
 import com.azure.core.http.netty.NettyAsyncHttpClientProvider;
 import com.azure.core.http.okhttp.OkHttpAsyncClientProvider;
 import com.azure.core.http.policy.HttpLogOptions;
@@ -25,6 +27,7 @@
 import com.azure.identity.DefaultAzureCredentialBuilder;
 import com.azure.identity.EnvironmentCredentialBuilder;
 import com.azure.storage.common.implementation.Constants;
+import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Schedulers;
 
 import java.io.ByteArrayOutputStream;
@@ -38,7 +41,10 @@
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 import java.util.Locale;
 import java.util.Random;
 import java.util.UUID;
@@ -389,4 +395,23 @@ public static TokenCredential getTokenCredential(InterceptorManager interceptorM
             return builder.build();
         }
     }
+
+    public static class CaptureAllRequestsHttpClient implements HttpClient {
+        private final List<HttpRequest> capturedRequests = new ArrayList<>();
+        private final HttpClient inner;
+
+        public CaptureAllRequestsHttpClient(HttpClient inner) {
+            this.inner = inner;
+        }
+
+        @Override
+        public Mono<HttpResponse> send(HttpRequest request) {
+            capturedRequests.add(request);
+            return inner.send(request);
+        }
+
+        public List<HttpRequest> getCapturedRequests() {
+            return Collections.unmodifiableList(capturedRequests);
+        }
+    }
 }