Fix MockPartialResponsePolicy and add comprehensive smart retry tests mirroring .NET implementation

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

@@ -23,9 +23,26 @@ public class MockPartialResponsePolicy implements HttpPipelinePolicy {
     static final HttpHeaderName RANGE_HEADER = HttpHeaderName.RANGE;
     private int tries;
     private final List<String> rangeHeaders = new ArrayList<>();
+    private final int maxBytesPerResponse;  // Maximum bytes to return before simulating timeout
 
+    /**
+     * Creates a MockPartialResponsePolicy that simulates network interruptions.
+     * 
+     * @param tries Number of times to simulate interruptions (0 = no interruptions)
+     */
     public MockPartialResponsePolicy(int tries) {
+        this(tries, 560);  // Default: return up to 560 bytes before interrupting (enough for 1 segment + header)
+    }
+
+    /**
+     * Creates a MockPartialResponsePolicy with configurable interruption behavior.
+     * 
+     * @param tries Number of times to simulate interruptions (0 = no interruptions)
+     * @param maxBytesPerResponse Maximum bytes to return in each interrupted response
+     */
+    public MockPartialResponsePolicy(int tries, int maxBytesPerResponse) {
         this.tries = tries;
+        this.maxBytesPerResponse = maxBytesPerResponse;
     }
 
     @Override
@@ -51,21 +68,53 @@ public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineN
                 return Mono.just(response);
             } else {
                 this.tries -= 1;
+                // 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. Simply propagate
-                        // the original response to avoid test failures when the service unexpectedly returns an
-                        // empty body (for example, after a retry on the underlying transport).
+                        // If no body was returned, don't attempt to slice a partial response
                         return Mono.just(response);
                     }
-                    ByteBuffer firstBuffer = bodyBuffers.get(0);
-                    byte firstByte = firstBuffer.get();
-
-                    // Simulate partial response by returning the first byte only from the requested range and timeout
+
+                    // 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(ByteBuffer.wrap(new byte[] { firstByte }))
-                                    .concatWith(Flux.error(new IOException("Simulated timeout")))
-                    ));
+                        Flux.just(limited).concatWith(Flux.error(new IOException("Simulated timeout")))));
                 });
             }
         });