Revert to collectList approach with proper byte limiting for structured messages

The streaming approach had issues with buffer consumption. Reverted to collecting
buffers first but with proper byte limiting (560 bytes default) to ensure at least
one segment completes before interruption. This properly tests smart retry from
segment boundaries.

Co-authored-by: gunjansingh-msft <[email protected]>

Author: Copilot

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

@@ -68,51 +68,55 @@ public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineN
                 return Mono.just(response);
             } else {
                 this.tries -= 1;
-                // Simulate partial response by limiting the amount of data returned from the stream
-                // before throwing an IOException to simulate a network interruption.
-                // This tests smart retry behavior where downloads should resume from the last
-                // complete segment boundary after each interruption.
-                Flux<ByteBuffer> interruptedBody = limitAndInterruptStream(response.getBody(), maxBytesPerResponse);
-                return Mono.just(new MockDownloadHttpResponse(response, 206, interruptedBody));
-            }
-        });
-    }
-
-    /**
-     * Limits a stream to return at most maxBytes before throwing an IOException.
-     */
-    private Flux<ByteBuffer> limitAndInterruptStream(Flux<ByteBuffer> body, int maxBytes) {
-        return Flux.defer(() -> {
-            final int[] bytesEmitted = new int[] {0};
-            return body.concatMap(buffer -> {
-                int remaining = maxBytes - bytesEmitted[0];
-                if (remaining <= 0) {
-                    // Already emitted enough bytes, throw error now
-                    return Flux.error(new IOException("Simulated timeout"));
-                }
-                
-                int bytesToEmit = Math.min(buffer.remaining(), remaining);
-                if (bytesToEmit < buffer.remaining()) {
-                    // Need to slice the buffer
-                    ByteBuffer limited = ByteBuffer.allocate(bytesToEmit);
-                    int originalLimit = buffer.limit();
-                    buffer.limit(buffer.position() + bytesToEmit);
-                    limited.put(buffer);
-                    buffer.limit(originalLimit);
-                    limited.flip();
-                    bytesEmitted[0] += bytesToEmit;
-                    // Emit the limited buffer, then error
-                    return Flux.just(limited).concatWith(Flux.error(new IOException("Simulated timeout")));
-                } else {
-                    // Emit the full buffer and continue
-                    bytesEmitted[0] += bytesToEmit;
-                    if (bytesEmitted[0] >= maxBytes) {
-                        // Reached the limit, emit this buffer then error
-                        return Flux.just(buffer).concatWith(Flux.error(new IOException("Simulated timeout")));
+                // Collect the body to be able to slice it properly
+                return response.getBody().collectList().flatMap(bodyBuffers -> {
+                    if (bodyBuffers.isEmpty()) {
+                        // If no body was returned, don't attempt to slice a partial response
+                        return Mono.just(response);
                     }
-                    return Flux.just(buffer);
-                }
-            });
+                    
+                    // Calculate total bytes available
+                    int totalBytes = bodyBuffers.stream().mapToInt(ByteBuffer::remaining).sum();
+                    
+                    // Determine how many bytes to return (limited by maxBytesPerResponse)
+                    int bytesToReturn = Math.min(totalBytes, maxBytesPerResponse);
+                    
+                    if (bytesToReturn >= totalBytes) {
+                        // Return all data and still throw error to simulate interruption during next chunk
+                        return Mono.just(new MockDownloadHttpResponse(response, 206,
+                            Flux.fromIterable(bodyBuffers)
+                                .concatWith(Flux.error(new IOException("Simulated timeout")))));
+                    }
+                    
+                    // Create a new buffer with limited bytes
+                    ByteBuffer limited = ByteBuffer.allocate(bytesToReturn);
+                    int bytesCollected = 0;
+                    
+                    for (ByteBuffer buffer : bodyBuffers) {
+                        int bufferRemaining = buffer.remaining();
+                        int bytesNeeded = bytesToReturn - bytesCollected;
+                        
+                        if (bufferRemaining <= bytesNeeded) {
+                            // Take the entire buffer
+                            limited.put(buffer);
+                            bytesCollected += bufferRemaining;
+                        } else {
+                            // Take only part of this buffer
+                            ByteBuffer slice = buffer.duplicate();
+                            slice.limit(slice.position() + bytesNeeded);
+                            limited.put(slice);
+                            bytesCollected += bytesNeeded;
+                            break;
+                        }
+                    }
+                    
+                    limited.flip();
+                    
+                    // Return the limited buffer and simulate timeout
+                    return Mono.just(new MockDownloadHttpResponse(response, 206,
+                        Flux.just(limited).concatWith(Flux.error(new IOException("Simulated timeout")))));
+                });
+            }
         });
     }