add test back

Author: GarrettBeatty

sdk/test/Services/S3/UnitTests/Custom/MultipartDownloadManagerTests.cs

@@ -1866,6 +1866,97 @@ public async Task StartDownloadsAsync_SinglePart_ReturnsImmediatelyWithoutBackgr
 
         #region Capacity Checking Tests
 
+        [TestMethod]
+        public async Task Discovery_PartStrategy_CallsWaitForCapacityAsync()
+        {
+            // Arrange - PART strategy should check capacity during discovery
+            var capacityCallCount = 0;
+            var mockDataHandler = new Mock<IPartDataHandler>();
+            
+            mockDataHandler
+                .Setup(x => x.WaitForCapacityAsync(It.IsAny<CancellationToken>()))
+                .Returns(() =>
+                {
+                    Interlocked.Increment(ref capacityCallCount);
+                    return Task.CompletedTask;
+                });
+            
+            mockDataHandler
+                .Setup(x => x.ProcessPartAsync(It.IsAny<int>(), It.IsAny<GetObjectResponse>(), It.IsAny<CancellationToken>()))
+                .Returns(Task.CompletedTask);
+            
+            var totalObjectSize = 24 * 1024 * 1024; // 24MB -> 3 parts @ 8MB
+            var partSize = 8 * 1024 * 1024; // 8MB
+            var totalParts = 3;
+            
+            var mockResponse = MultipartDownloadTestHelpers.CreateMultipartFirstPartResponse(
+                partSize, totalParts, totalObjectSize, "test-etag");
+            
+            var mockClient = MultipartDownloadTestHelpers.CreateMockS3Client(
+                (req, ct) => Task.FromResult(mockResponse));
+            
+            var request = MultipartDownloadTestHelpers.CreateOpenStreamRequest(
+                downloadType: MultipartDownloadType.PART);
+            var config = MultipartDownloadTestHelpers.CreateBufferedDownloadConfiguration();
+            var coordinator = new MultipartDownloadManager(
+                mockClient.Object, request, config, mockDataHandler.Object);
+
+            // Act
+            var result = await coordinator.StartDownloadAsync(null, CancellationToken.None);
+
+            // Assert
+            Assert.IsNotNull(result);
+            Assert.AreEqual(3, result.TotalParts);
+            Assert.AreEqual(1, capacityCallCount, 
+                "PART strategy should call WaitForCapacityAsync during Part 1 discovery");
+        }
+
+
+        [TestMethod]
+        public async Task Discovery_RangeStrategy_CallsWaitForCapacityAsync()
+        {
+            // Arrange - RANGE strategy should also check capacity during discovery
+            var capacityCallCount = 0;
+            var mockDataHandler = new Mock<IPartDataHandler>();
+            
+            mockDataHandler
+                .Setup(x => x.WaitForCapacityAsync(It.IsAny<CancellationToken>()))
+                .Returns(() =>
+                {
+                    Interlocked.Increment(ref capacityCallCount);
+                    return Task.CompletedTask;
+                });
+            
+            mockDataHandler
+                .Setup(x => x.ProcessPartAsync(It.IsAny<int>(), It.IsAny<GetObjectResponse>(), It.IsAny<CancellationToken>()))
+                .Returns(Task.CompletedTask);
+            
+            var totalObjectSize = 17 * 1024 * 1024; // 17MB -> 3 parts @ 8MB
+            var partSize = 8 * 1024 * 1024; // 8MB
+            
+            var mockResponse = MultipartDownloadTestHelpers.CreateRangeResponse(
+                0, partSize - 1, totalObjectSize, "test-etag");
+            
+            var mockClient = MultipartDownloadTestHelpers.CreateMockS3Client(
+                (req, ct) => Task.FromResult(mockResponse));
+            
+            var request = MultipartDownloadTestHelpers.CreateOpenStreamRequest(
+                partSize: partSize,
+                downloadType: MultipartDownloadType.RANGE);
+            var config = MultipartDownloadTestHelpers.CreateBufferedDownloadConfiguration();
+            var coordinator = new MultipartDownloadManager(
+                mockClient.Object, request, config, mockDataHandler.Object);
+
+            // Act
+            var result = await coordinator.StartDownloadAsync(null, CancellationToken.None);
+
+            // Assert
+            Assert.IsNotNull(result);
+            Assert.AreEqual(3, result.TotalParts); // 17MB / 8MB = 3 parts (ceiling)
+            Assert.AreEqual(1, capacityCallCount, 
+                "RANGE strategy should call WaitForCapacityAsync during Part 1 discovery");
+        }
+
         [TestMethod]
         public async Task MultipleDownloads_WithSharedHttpThrottler_RespectsLimits()
         {
@@ -2077,50 +2168,6 @@ public async Task Discovery_SinglePart_StillCallsCapacityCheck()
                 "Even single-part downloads should call WaitForCapacityAsync during discovery");
         }
 
-        [TestMethod]
-        public async Task Discovery_RangeStrategy_CallsWaitForCapacityAsync()
-        {
-            // Arrange - RANGE strategy should also check capacity during discovery
-            var capacityCallCount = 0;
-            var mockDataHandler = new Mock<IPartDataHandler>();
-            
-            mockDataHandler
-                .Setup(x => x.WaitForCapacityAsync(It.IsAny<CancellationToken>()))
-                .Returns(() =>
-                {
-                    Interlocked.Increment(ref capacityCallCount);
-                    return Task.CompletedTask;
-                });
-            
-            mockDataHandler
-                .Setup(x => x.ProcessPartAsync(It.IsAny<int>(), It.IsAny<GetObjectResponse>(), It.IsAny<CancellationToken>()))
-                .Returns(Task.CompletedTask);
-            
-            var totalObjectSize = 17 * 1024 * 1024; // 17MB -> 3 parts @ 8MB
-            var partSize = 8 * 1024 * 1024; // 8MB
-            
-            var mockResponse = MultipartDownloadTestHelpers.CreateRangeResponse(
-                0, partSize - 1, totalObjectSize, "test-etag");
-            
-            var mockClient = MultipartDownloadTestHelpers.CreateMockS3Client(
-                (req, ct) => Task.FromResult(mockResponse));
-            
-            var request = MultipartDownloadTestHelpers.CreateOpenStreamRequest(
-                partSize: partSize,
-                downloadType: MultipartDownloadType.RANGE);
-            var config = MultipartDownloadTestHelpers.CreateBufferedDownloadConfiguration();
-            var coordinator = new MultipartDownloadManager(
-                mockClient.Object, request, config, mockDataHandler.Object);
-
-            // Act
-            var result = await coordinator.StartDownloadAsync(null, CancellationToken.None);
-
-            // Assert
-            Assert.IsNotNull(result);
-            Assert.AreEqual(3, result.TotalParts); // 17MB / 8MB = 3 parts (ceiling)
-            Assert.AreEqual(1, capacityCallCount, 
-                "RANGE strategy should call WaitForCapacityAsync during Part 1 discovery");
-        }
 
         #endregion