Fix for flakey Download 304 tests to check LastModified time from Upload (#33325)

Author: amnguye

sdk/storage/Azure.Storage.Blobs/tests/BlobBaseClientTests.cs

@@ -1547,11 +1547,14 @@ public async Task DownloadTo_Initial304()
             // Upload a blob
             var data = GetRandomBuffer(Constants.KB);
             BlobClient blob = InstrumentClient(test.Container.GetBlobClient(GetNewBlobName()));
+            Response<BlobContentInfo> uploadResponse;
             using (var stream = new MemoryStream(data))
             {
-                await blob.UploadAsync(stream);
+                uploadResponse = await blob.UploadAsync(stream);
             }
 
+            CheckModifiedSinceAndWait(uploadResponse);
+
             // Add conditions to cause a failure and ensure we don't explode
             Response result = await blob.DownloadToAsync(
                 Stream.Null,
@@ -1571,11 +1574,14 @@ public async Task DownloadContent_Initial304()
             // Upload a blob
             var data = GetRandomBuffer(Constants.KB);
             BlobClient blob = InstrumentClient(test.Container.GetBlobClient(GetNewBlobName()));
+            Response<BlobContentInfo> uploadResponse;
             using (var stream = new MemoryStream(data))
             {
-                await blob.UploadAsync(stream);
+                uploadResponse = await blob.UploadAsync(stream);
             }
 
+            CheckModifiedSinceAndWait(uploadResponse);
+
             // Add conditions to cause a failure and ensure we don't explode
             Response<BlobDownloadResult> result = await blob.DownloadContentAsync(new BlobDownloadOptions
             {
@@ -1596,11 +1602,14 @@ public async Task DownloadStreaming_Initial304()
             // Upload a blob
             var data = GetRandomBuffer(Constants.KB);
             BlobClient blob = InstrumentClient(test.Container.GetBlobClient(GetNewBlobName()));
+            Response<BlobContentInfo> uploadResponse;
             using (var stream = new MemoryStream(data))
             {
-                await blob.UploadAsync(stream);
+                uploadResponse = await blob.UploadAsync(stream);
             }
 
+            CheckModifiedSinceAndWait(uploadResponse);
+
             // Add conditions to cause a failure and ensure we don't explode
             Response<BlobDownloadStreamingResult> result = await blob.DownloadStreamingAsync(new BlobDownloadOptions
             {
@@ -1621,11 +1630,14 @@ public async Task Download_Initial304()
             // Upload a blob
             var data = GetRandomBuffer(Constants.KB);
             BlobClient blob = InstrumentClient(test.Container.GetBlobClient(GetNewBlobName()));
+            Response<BlobContentInfo> uploadResponse;
             using (var stream = new MemoryStream(data))
             {
-                await blob.UploadAsync(stream);
+                uploadResponse = await blob.UploadAsync(stream);
             }
 
+            CheckModifiedSinceAndWait(uploadResponse);
+
             // Add conditions to cause a failure and ensure we don't explode
             Response<BlobDownloadInfo> result = await blob.DownloadAsync(
                 conditions: new BlobRequestConditions

sdk/storage/Azure.Storage.Blobs/tests/BlobTestBase.cs

@@ -6,6 +6,7 @@
 using System.IO;
 using System.Net;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
 using Azure.Core.TestFramework;
@@ -573,5 +574,25 @@ public void Handle(BlobQueryError blobQueryError)
                 Assert.AreEqual(_expectedBlobQueryError.Position, blobQueryError.Position);
             }
         }
+
+        /// <summary>
+        /// For our Download Initial Response 304 tests, sometimes the Modfified Since time
+        /// comes back to be later than the actual Utc time in the environment that ran the test.
+        /// We run this method to check the IfModified time and if it's later than the time, then
+        /// let's wait until that time has past and then continue.
+        ///
+        /// Or else when we attempt to the Download at set the LastModified time before that, it will
+        /// send us back a 200, or even an error for sending a LastModified time that's in the future.
+        ///
+        /// At most it's a 1 second difference wait. But most of the time there's no wait.
+        /// </summary>
+        public void CheckModifiedSinceAndWait(Response<BlobContentInfo> uploadResponse)
+        {
+            DateTimeOffset offsetNow = Recording.UtcNow;
+            if (DateTimeOffset.Compare(uploadResponse.Value.LastModified, offsetNow) > 0)
+            {
+                Thread.Sleep(uploadResponse.Value.LastModified.Subtract(offsetNow));
+            }
+        }
     }
 }