[release/17.0] Let REPL test readiness wait cover a cold container image pull

tests/Aspire.Hosting.MongoDB.Tests/Aspire.Hosting.MongoDB.Tests.csproj

@@ -5,6 +5,7 @@
   </PropertyGroup>
   <ItemGroup>
     <Compile Include="..\Shared\ContainerReplCommandTestBase.cs" Link="ContainerReplCommandTestBase.cs" />
+    <Compile Include="$(TestsSharedDir)AsyncTestHelpers.cs" Link="shared/AsyncTestHelpers.cs" />
   </ItemGroup>
   <ItemGroup>
     <ProjectReference Include="..\..\src\Aspire.Hosting.MongoDB\Aspire.Hosting.MongoDB.csproj" />

tests/Aspire.Hosting.PostgreSQL.Tests/Aspire.Hosting.PostgreSQL.Tests.csproj

@@ -16,6 +16,7 @@
     <ProjectReference Include="..\..\src\Aspire.Hosting.PostgreSQL\Aspire.Hosting.PostgreSQL.csproj" />
     <ProjectReference Include="..\..\src\Components\Aspire.Npgsql\Aspire.Npgsql.csproj" />
     <ProjectReference Include="..\Aspire.Hosting.TestUtilities\Aspire.Hosting.TestUtilities.csproj" />
+    <PackageReference Include="Microsoft.Extensions.TimeProvider.Testing" />
   </ItemGroup>
 
   <ItemGroup>

tests/Aspire.Hosting.PostgreSQL.Tests/PostgresReplCommandTests.cs

@@ -6,6 +6,8 @@
 using Aspire.Hosting.Tests;
 using Aspire.Hosting.Utils;
 using Aspire.TestUtilities;
+using Microsoft.AspNetCore.InternalTesting;
+using Microsoft.Extensions.Time.Testing;
 
 #pragma warning disable ASPIRETERMINAL001
 
@@ -72,6 +74,28 @@ public async Task ReplRejectsMissingTargetPort()
         Assert.Equal("The PostgreSQL REPL port is not available.", exception.Message);
     }
 
+    [Fact]
+    public async Task ReplReadinessWaitCoversColdImagePull()
+    {
+        // A cold postgres image pull took ~55s in CI before the container could even start, so the
+        // readiness wait must still be alive well past that point.
+        var coldImagePullAndStartup = TimeSpan.FromMinutes(2);
+
+        using var builder = TestDistributedApplicationBuilder.Create();
+        var postgres = builder.AddPostgres("postgres");
+        await using var app = builder.Build();
+        var timeProvider = new FakeTimeProvider();
+
+        var waitTask = WaitForReplResourceHealthyAsync(app, postgres.Resource.Name, timeProvider);
+        timeProvider.Advance(coldImagePullAndStartup);
+        Assert.False(waitTask.IsCompleted);
+
+        // The wait only observes the resource once it reaches a terminal state, so the wait failing for
+        // that reason (rather than being cancelled) shows the budget had not run out.
+        await app.ResourceNotifications.PublishUpdateAsync(postgres.Resource, s => s with { State = KnownResourceStates.FailedToStart });
+        await Assert.ThrowsAsync<DistributedApplicationException>(() => waitTask).DefaultTimeout();
+    }
+
     [Theory]
     [InlineData(null)]
     [InlineData(5433)]

tests/Aspire.Hosting.Redis.Tests/Aspire.Hosting.Redis.Tests.csproj

@@ -2,6 +2,7 @@
 
   <ItemGroup>
     <Compile Include="..\Shared\ContainerReplCommandTestBase.cs" Link="ContainerReplCommandTestBase.cs" />
+    <Compile Include="$(TestsSharedDir)AsyncTestHelpers.cs" Link="shared/AsyncTestHelpers.cs" />
   </ItemGroup>
 
   <PropertyGroup>

tests/Aspire.Hosting.SqlServer.Tests/Aspire.Hosting.SqlServer.Tests.csproj

@@ -7,6 +7,7 @@
   </PropertyGroup>
   <ItemGroup>
     <Compile Include="..\Shared\ContainerReplCommandTestBase.cs" Link="ContainerReplCommandTestBase.cs" />
+    <Compile Include="$(TestsSharedDir)AsyncTestHelpers.cs" Link="shared/AsyncTestHelpers.cs" />
   </ItemGroup>
   <ItemGroup>
     <ProjectReference Include="..\..\src\Aspire.Hosting.SqlServer\Aspire.Hosting.SqlServer.csproj" />

tests/Shared/ContainerReplCommandTestBase.cs

@@ -5,6 +5,7 @@
 using Aspire.Hosting.ApplicationModel;
 using Aspire.Hosting.Tests.Utils;
 using Aspire.Hosting.Utils;
+using Microsoft.AspNetCore.InternalTesting;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging.Abstractions;
 
@@ -160,10 +161,7 @@ protected static async Task VerifyReplAsync(
             await app.StartAsync(startup.Token);
         }
 
-        using (var readiness = new CancellationTokenSource(TimeSpan.FromMinutes(1)))
-        {
-            await app.ResourceNotifications.WaitForResourceHealthyAsync(resource.Name, readiness.Token);
-        }
+        await WaitForReplResourceHealthyAsync(app, resource.Name, TimeProvider.System);
 
         using var cancellation = new CancellationTokenSource(TimeSpan.FromMinutes(1));
         await using var terminal = await TerminalCommandTestHelpers.ExecuteTerminalCommandAsync(app, resource, "repl", cancellation.Token);
@@ -187,4 +185,13 @@ protected static async Task VerifyReplAsync(
             throw new TimeoutException($"REPL did not produce the expected output. Terminal screen:{Environment.NewLine}{terminal.GetScreenText()}", exception);
         }
     }
+
+    protected static async Task WaitForReplResourceHealthyAsync(DistributedApplication app, string resourceName, TimeProvider timeProvider)
+    {
+        // StartAsync returns before DCP pulls the container image, so this wait also has to cover a
+        // cold image pull. A cold postgres pull alone took ~55s in CI, which exhausted the previous
+        // 1-minute budget before the container could become healthy.
+        using var readiness = new CancellationTokenSource(TestConstants.ExtraLongTimeoutTimeSpan, timeProvider);
+        await app.ResourceNotifications.WaitForResourceHealthyAsync(resourceName, readiness.Token);
+    }
 }