diff --git a/test/Libraries/Microsoft.Extensions.AI.Integration.Tests/ImageGeneratorIntegrationTests.cs b/test/Libraries/Microsoft.Extensions.AI.Integration.Tests/ImageGeneratorIntegrationTests.cs
index 8e3078efbb5..76b08941bc5 100644
--- a/test/Libraries/Microsoft.Extensions.AI.Integration.Tests/ImageGeneratorIntegrationTests.cs
+++ b/test/Libraries/Microsoft.Extensions.AI.Integration.Tests/ImageGeneratorIntegrationTests.cs
@@ -43,13 +43,23 @@ public virtual async Task GenerateImagesAsync_SingleImageGeneration()
 
         Assert.NotNull(response);
         Assert.NotEmpty(response.Contents);
-        Assert.Single(response.Contents);
 
-        var content = response.Contents[0];
-        Assert.IsType<DataContent>(content);
-        var dataContent = (DataContent)content;
-        Assert.False(dataContent.Data.IsEmpty);
-        Assert.StartsWith("image/", dataContent.MediaType, StringComparison.Ordinal);
+        var content = Assert.Single(response.Contents);
+        switch (content)
+        {
+            case UriContent uc:
+                Assert.StartsWith("http", uc.Uri.Scheme, StringComparison.Ordinal);
+                break;
+
+            case DataContent dc:
+                Assert.False(dc.Data.IsEmpty);
+                Assert.StartsWith("image/", dc.MediaType, StringComparison.Ordinal);
+                break;
+
+            default:
+                Assert.Fail($"Unexpected content type: {content.GetType()}");
+                break;
+        }
     }
 
     [ConditionalFact]