diff --git a/src/Aspire.Hosting.MongoDB/MongoDBBuilderExtensions.cs b/src/Aspire.Hosting.MongoDB/MongoDBBuilderExtensions.cs index 421dccd19a..faa1bee4b4 100644 --- a/src/Aspire.Hosting.MongoDB/MongoDBBuilderExtensions.cs +++ b/src/Aspire.Hosting.MongoDB/MongoDBBuilderExtensions.cs @@ -1,6 +1,7 @@ // Licensed to the .NET Foundation under one or more agreements. // The .NET Foundation licenses this file to you under the MIT license. +using System.Text; using Aspire.Hosting.ApplicationModel; using Aspire.Hosting.MongoDB; using Aspire.Hosting.Utils; @@ -94,7 +95,8 @@ public static IResourceBuilder WithMongoExpress(this IResourceBuilder b .WithImageRegistry(MongoDBContainerImageTags.MongoExpressRegistry) .WithEnvironment(context => ConfigureMongoExpressContainer(context, builder.Resource)) .WithHttpEndpoint(targetPort: 8081, name: "http") - .ExcludeFromManifest(); + .ExcludeFromManifest() + .WaitFor(builder); configureContainer?.Invoke(resourceBuilder); @@ -161,11 +163,96 @@ public static IResourceBuilder WithInitBindMount(this IRe return builder.WithBindMount(source, "/docker-entrypoint-initdb.d", isReadOnly); } + /// + /// Adds a replica set to the MongoDB server resource. + /// + /// The MongoDB server resource. + /// The name of the replica set. If not provided, defaults to rs0. + /// A reference to the . + public static IResourceBuilder WithReplicaSet(this IResourceBuilder builder, string? replicaSetName = null) + { + if (builder.Resource.TryGetLastAnnotation(out _)) + { + throw new InvalidOperationException("A replica set has already been added to the MongoDB server resource."); + } + + replicaSetName ??= "rs0"; + + var port = SetPortAndTargetToBeSame(builder); + + // Add a container that initializes the replica set + var init = builder.ApplicationBuilder + .AddDockerfile("replicaset-init", GetReplicaSetInitDockerfileDir(replicaSetName, builder.Resource.Name, port)) + + // We don't want to wait for the healthchecks to be successful since the initialization is required for that. However, we also don't want this to start + // up until the database itself is ready + .WaitFor(builder, includeHealthChecks: false); + + return builder + .WithAnnotation(new MongoDbReplicaSetAnnotation(replicaSetName, init)) + .WithArgs("--replSet", replicaSetName, "--bind_ip_all", "--port", $"{port}"); + + static int SetPortAndTargetToBeSame(IResourceBuilder builder) + { + foreach (var endpoint in builder.Resource.Annotations.OfType()) + { + if (endpoint.Name == MongoDBServerResource.PrimaryEndpointName) + { + if (endpoint.Port is { } port) + { + endpoint.TargetPort = port; + } + + if (endpoint.TargetPort is not { } targetPort) + { + throw new InvalidOperationException("Target port is not set."); + } + + // In the case of replica sets, the port and target port should be the same and is not proxied + endpoint.IsProxied = false; + + return targetPort; + } + } + + throw new InvalidOperationException("No endpoint found for the MongoDB server resource."); + } + + // See the conversation about setting up replica sets in Docker here: https://github.com/docker-library/mongo/issues/246 + static string GetReplicaSetInitDockerfileDir(string replicaSet, string host, int port) + { + var dir = Directory.CreateTempSubdirectory("aspire.mongo").FullName; + + var rsInitContents = $$"""rs.initiate({ _id:'{{replicaSet}}', members:[{_id:0,host:'localhost:{{port}}'}]})"""; + var init = Path.Combine(dir, "rs.js"); + File.WriteAllText(init, rsInitContents); + + var dockerfile = Path.Combine(dir, "Dockerfile"); + File.WriteAllText(dockerfile, $""" + FROM {MongoDBContainerImageTags.Image}:{MongoDBContainerImageTags.Tag} + WORKDIR /rsinit + ADD rs.js rs.js + ENTRYPOINT ["mongosh", "--port", "{port}", "--host", "{host}", "rs.js"] + """); + return dir; + } + } + private static void ConfigureMongoExpressContainer(EnvironmentCallbackContext context, MongoDBServerResource resource) { - // Mongo Exporess assumes Mongo is being accessed over a default Aspire container network and hardcodes the resource address + var sb = new StringBuilder($"mongodb://{resource.Name}:{resource.PrimaryEndpoint.TargetPort}/?directConnection=true"); + + if (resource.TryGetLastAnnotation(out var replica)) + { + sb.Append('&'); + sb.Append(MongoDbReplicaSetAnnotation.QueryName); + sb.Append('='); + sb.Append(replica.ReplicaSetName); + } + + // Mongo Express assumes Mongo is being accessed over a default Aspire container network and hardcodes the resource address // This will need to be refactored once updated service discovery APIs are available - context.EnvironmentVariables.Add("ME_CONFIG_MONGODB_URL", $"mongodb://{resource.Name}:{resource.PrimaryEndpoint.TargetPort}/?directConnection=true"); + context.EnvironmentVariables.Add("ME_CONFIG_MONGODB_URL", sb.ToString()); context.EnvironmentVariables.Add("ME_CONFIG_BASICAUTH", "false"); } } diff --git a/src/Aspire.Hosting.MongoDB/MongoDBDatabaseResource.cs b/src/Aspire.Hosting.MongoDB/MongoDBDatabaseResource.cs index b41572913f..5a80dea06c 100644 --- a/src/Aspire.Hosting.MongoDB/MongoDBDatabaseResource.cs +++ b/src/Aspire.Hosting.MongoDB/MongoDBDatabaseResource.cs @@ -27,7 +27,17 @@ public MongoDBDatabaseResource(string name, string databaseName, MongoDBServerRe /// Gets the connection string expression for the MongoDB database. /// public ReferenceExpression ConnectionStringExpression - => ReferenceExpression.Create($"{Parent}/{DatabaseName}"); + { + get + { + var builder = new ReferenceExpressionBuilder(); + + Parent.AppendConnectionString(builder); + Parent.AppendSuffix(builder, DatabaseName); + + return builder.Build(); + } + } /// /// Gets the parent MongoDB container resource. diff --git a/src/Aspire.Hosting.MongoDB/MongoDBServerResource.cs b/src/Aspire.Hosting.MongoDB/MongoDBServerResource.cs index abb4298320..5312d3b22c 100644 --- a/src/Aspire.Hosting.MongoDB/MongoDBServerResource.cs +++ b/src/Aspire.Hosting.MongoDB/MongoDBServerResource.cs @@ -21,9 +21,59 @@ public class MongoDBServerResource(string name) : ContainerResource(name), IReso /// /// Gets the connection string for the MongoDB server. /// - public ReferenceExpression ConnectionStringExpression => - ReferenceExpression.Create( - $"mongodb://{PrimaryEndpoint.Property(EndpointProperty.Host)}:{PrimaryEndpoint.Property(EndpointProperty.Port)}"); + public ReferenceExpression ConnectionStringExpression + { + get + { + var builder = new ReferenceExpressionBuilder(); + + AppendConnectionString(builder); + AppendSuffix(builder); + + return builder.Build(); + } + } + + internal void AppendConnectionString(ReferenceExpressionBuilder builder) + { + builder.AppendLiteral("mongodb://"); + builder.AppendFormatted(PrimaryEndpoint.Property(EndpointProperty.Host)); + builder.AppendLiteral(":"); + builder.AppendFormatted(PrimaryEndpoint.Property(EndpointProperty.Port)); + } + + /// + /// Handles adding the rest of the connection string. + /// + /// - If a database name is provided, it will be appended to the connection string. + /// - If a replica set name is provided, it will be appended to the connection string. + /// - If no database but a replica set is provided, a '/' must be inserted before the '?' + /// + internal bool AppendSuffix(ReferenceExpressionBuilder builder, string? dbName = null) + { + if (dbName is { }) + { + builder.AppendLiteral("/"); + builder.AppendFormatted(dbName); + } + + if (Annotations.OfType().FirstOrDefault() is { ReplicaSetName: { } replicaSetName }) + { + if (dbName is null) + { + builder.AppendLiteral("/"); + } + + builder.AppendLiteral("?"); + builder.AppendLiteral(MongoDbReplicaSetAnnotation.QueryName); + builder.AppendLiteral("="); + builder.AppendLiteral(replicaSetName); + + return true; + } + + return false; + } private readonly Dictionary _databases = new Dictionary(StringComparers.ResourceName); diff --git a/src/Aspire.Hosting.MongoDB/MongoDbReplicaSetAnnotation.cs b/src/Aspire.Hosting.MongoDB/MongoDbReplicaSetAnnotation.cs new file mode 100644 index 0000000000..3c8bc1145a --- /dev/null +++ b/src/Aspire.Hosting.MongoDB/MongoDbReplicaSetAnnotation.cs @@ -0,0 +1,9 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +namespace Aspire.Hosting.ApplicationModel; + +internal sealed record MongoDbReplicaSetAnnotation(string ReplicaSetName, IResourceBuilder InitContainer) : IResourceAnnotation +{ + internal const string QueryName = "replicaSet"; +} diff --git a/src/Aspire.Hosting.MongoDB/PublicAPI.Unshipped.txt b/src/Aspire.Hosting.MongoDB/PublicAPI.Unshipped.txt index 074c6ad103..336a347488 100644 --- a/src/Aspire.Hosting.MongoDB/PublicAPI.Unshipped.txt +++ b/src/Aspire.Hosting.MongoDB/PublicAPI.Unshipped.txt @@ -1,2 +1,2 @@ #nullable enable - +static Aspire.Hosting.MongoDBBuilderExtensions.WithReplicaSet(this Aspire.Hosting.ApplicationModel.IResourceBuilder! builder, string? replicaSetName = null) -> Aspire.Hosting.ApplicationModel.IResourceBuilder! diff --git a/src/Aspire.Hosting/PublicAPI.Unshipped.txt b/src/Aspire.Hosting/PublicAPI.Unshipped.txt index 964f8126fc..9efe90ef67 100644 --- a/src/Aspire.Hosting/PublicAPI.Unshipped.txt +++ b/src/Aspire.Hosting/PublicAPI.Unshipped.txt @@ -86,6 +86,7 @@ Aspire.Hosting.DistributedApplicationExecutionContextOptions.Operation.get -> As Aspire.Hosting.DistributedApplicationExecutionContextOptions.ServiceProvider.get -> System.IServiceProvider? Aspire.Hosting.DistributedApplicationExecutionContextOptions.ServiceProvider.set -> void static Aspire.Hosting.ResourceBuilderExtensions.WaitFor(this Aspire.Hosting.ApplicationModel.IResourceBuilder! builder, Aspire.Hosting.ApplicationModel.IResourceBuilder! dependency) -> Aspire.Hosting.ApplicationModel.IResourceBuilder! +static Aspire.Hosting.ResourceBuilderExtensions.WaitFor(this Aspire.Hosting.ApplicationModel.IResourceBuilder! builder, Aspire.Hosting.ApplicationModel.IResourceBuilder! dependency, bool includeHealthChecks) -> Aspire.Hosting.ApplicationModel.IResourceBuilder! static Aspire.Hosting.ResourceBuilderExtensions.WaitForCompletion(this Aspire.Hosting.ApplicationModel.IResourceBuilder! builder, Aspire.Hosting.ApplicationModel.IResourceBuilder! dependency, int exitCode = 0) -> Aspire.Hosting.ApplicationModel.IResourceBuilder! static Aspire.Hosting.ResourceBuilderExtensions.WithHealthCheck(this Aspire.Hosting.ApplicationModel.IResourceBuilder! builder, string! key) -> Aspire.Hosting.ApplicationModel.IResourceBuilder! static readonly Aspire.Hosting.ApplicationModel.KnownResourceStates.Exited -> string! diff --git a/src/Aspire.Hosting/ResourceBuilderExtensions.cs b/src/Aspire.Hosting/ResourceBuilderExtensions.cs index 72804b81ad..05157f5eb7 100644 --- a/src/Aspire.Hosting/ResourceBuilderExtensions.cs +++ b/src/Aspire.Hosting/ResourceBuilderExtensions.cs @@ -596,6 +596,37 @@ public static IResourceBuilder ExcludeFromManifest(this IResourceBuilder /// public static IResourceBuilder WaitFor(this IResourceBuilder builder, IResourceBuilder dependency) where T : IResource + => builder.WaitFor(dependency, includeHealthChecks: true); + + /// + /// Waits for the dependency resource to enter the Running state before starting the resource. + /// + /// The type of the resource. + /// The resource builder for the resource that will be waiting. + /// The resource builder for the dependency resource. + /// Optionally includes checking the health checks. + /// The resource builder. + /// + /// This method is useful when a resource should wait until another has started running. This can help + /// reduce errors in logs during local development where dependency resources. + /// Some resources automatically register health checks with the application host container. For these + /// resources, calling also results + /// in the resource being blocked from starting until the health checks associated with the dependency resource + /// return . + /// The method can be used to associate + /// additional health checks with a resource. + /// + /// + /// Start message queue before starting the worker service. + /// + /// var builder = DistributedApplication.CreateBuilder(args); + /// var messaging = builder.AddRabbitMQ("messaging"); + /// builder.AddProject<Projects.MyApp>("myapp") + /// .WithReference(messaging) + /// .WaitFor(messaging); + /// + /// + public static IResourceBuilder WaitFor(this IResourceBuilder builder, IResourceBuilder dependency, bool includeHealthChecks) where T : IResource { builder.ApplicationBuilder.Eventing.Subscribe(builder.Resource, async (e, ct) => { @@ -632,7 +663,7 @@ public static IResourceBuilder WaitFor(this IResourceBuilder builder, I // If our dependency resource has health check annotations we want to wait until they turn healthy // otherwise we don't care about their health status. - if (dependency.Resource.TryGetAnnotationsOfType(out var _)) + if (includeHealthChecks && dependency.Resource.TryGetAnnotationsOfType(out var _)) { resourceLogger.LogInformation("Waiting for resource '{Name}' to become healthy.", dependency.Resource.Name); await rns.WaitForResourceAsync(dependency.Resource.Name, re => re.Snapshot.HealthStatus == HealthStatus.Healthy, cancellationToken: ct).ConfigureAwait(false); diff --git a/tests/Aspire.Hosting.MongoDB.Tests/AddMongoDBTests.cs b/tests/Aspire.Hosting.MongoDB.Tests/AddMongoDBTests.cs index afae345db3..23f4f11870 100644 --- a/tests/Aspire.Hosting.MongoDB.Tests/AddMongoDBTests.cs +++ b/tests/Aspire.Hosting.MongoDB.Tests/AddMongoDBTests.cs @@ -91,7 +91,33 @@ public async Task MongoDBCreatesConnectionString() Assert.Equal("mongodb://localhost:27017", await serverResource.GetConnectionStringAsync()); Assert.Equal("mongodb://{mongodb.bindings.tcp.host}:{mongodb.bindings.tcp.port}", serverResource.ConnectionStringExpression.ValueExpression); Assert.Equal("mongodb://localhost:27017/mydatabase", connectionString); - Assert.Equal("{mongodb.connectionString}/mydatabase", connectionStringResource.ConnectionStringExpression.ValueExpression); + Assert.Equal("mongodb://{mongodb.bindings.tcp.host}:{mongodb.bindings.tcp.port}/mydatabase", connectionStringResource.ConnectionStringExpression.ValueExpression); + } + + [Fact] + public async Task MongoDBCreatesConnectionStringWithReplicaSet() + { + var appBuilder = DistributedApplication.CreateBuilder(); + appBuilder + .AddMongoDB("mongodb") + .WithEndpoint("tcp", e => e.AllocatedEndpoint = new AllocatedEndpoint(e, "localhost", 27017)) + .WithReplicaSet("myreplset") + .AddDatabase("mydatabase"); + + using var app = appBuilder.Build(); + + var appModel = app.Services.GetRequiredService(); + + var dbResource = Assert.Single(appModel.Resources.OfType()); + var serverResource = dbResource.Parent as IResourceWithConnectionString; + var connectionStringResource = dbResource as IResourceWithConnectionString; + Assert.NotNull(connectionStringResource); + var connectionString = await connectionStringResource.GetConnectionStringAsync(); + + Assert.Equal("mongodb://localhost:27017/?replicaSet=myreplset", await serverResource.GetConnectionStringAsync()); + Assert.Equal("mongodb://{mongodb.bindings.tcp.host}:{mongodb.bindings.tcp.port}/?replicaSet=myreplset", serverResource.ConnectionStringExpression.ValueExpression); + Assert.Equal("mongodb://localhost:27017/mydatabase?replicaSet=myreplset", connectionString); + Assert.Equal("mongodb://{mongodb.bindings.tcp.host}:{mongodb.bindings.tcp.port}/mydatabase?replicaSet=myreplset", connectionStringResource.ConnectionStringExpression.ValueExpression); } [Fact] @@ -161,6 +187,31 @@ public async Task WithMongoExpressUsesContainerHost() }); } + [Fact] + public async Task WithMongoExpressWithReplicaSet() + { + using var builder = TestDistributedApplicationBuilder.Create(); + builder.AddMongoDB("mongo") + .WithReplicaSet("myreplicaset") + .WithMongoExpress(); + + var mongoExpress = Assert.Single(builder.Resources.OfType()); + + var env = await EnvironmentVariableEvaluator.GetEnvironmentVariablesAsync(mongoExpress, DistributedApplicationOperation.Run, TestServiceProvider.Instance); + + Assert.Collection(env, + e => + { + Assert.Equal("ME_CONFIG_MONGODB_URL", e.Key); + Assert.Equal($"mongodb://mongo:27017/?directConnection=true&replicaSet=myreplicaset", e.Value); + }, + e => + { + Assert.Equal("ME_CONFIG_BASICAUTH", e.Key); + Assert.Equal("false", e.Value); + }); + } + [Fact] public void WithMongoExpressOnMultipleResources() { @@ -201,7 +252,53 @@ public async Task VerifyManifest() expectedManifest = """ { "type": "value.v0", - "connectionString": "{mongo.connectionString}/mydb" + "connectionString": "mongodb://{mongo.bindings.tcp.host}:{mongo.bindings.tcp.port}/mydb" + } + """; + Assert.Equal(expectedManifest, dbManifest.ToString()); + } + + [InlineData(null)] + [InlineData(10002)] + [Theory] + public async Task VerifyManifestWithReplicaSet(int? customPort) + { + var appBuilder = DistributedApplication.CreateBuilder(); + var mongo = appBuilder.AddMongoDB("mongo", customPort) + .WithReplicaSet(); + var db = mongo.AddDatabase("mydb"); + + var mongoManifest = await ManifestUtils.GetManifest(mongo.Resource); + var dbManifest = await ManifestUtils.GetManifest(db.Resource); + + var expectedManifest = $$""" + { + "type": "container.v0", + "connectionString": "mongodb://{mongo.bindings.tcp.host}:{mongo.bindings.tcp.port}/?replicaSet=rs0", + "image": "{{MongoDBContainerImageTags.Registry}}/{{MongoDBContainerImageTags.Image}}:{{MongoDBContainerImageTags.Tag}}", + "args": [ + "--replSet", + "rs0", + "--bind_ip_all", + "--port", + "{{customPort ?? 27017}}" + ], + "bindings": { + "tcp": { + "scheme": "tcp", + "protocol": "tcp", + "transport": "tcp", + "targetPort": {{customPort ?? 27017}} + } + } + } + """; + Assert.Equal(expectedManifest, mongoManifest.ToString()); + + expectedManifest = """ + { + "type": "value.v0", + "connectionString": "mongodb://{mongo.bindings.tcp.host}:{mongo.bindings.tcp.port}/mydb?replicaSet=rs0" } """; Assert.Equal(expectedManifest, dbManifest.ToString()); @@ -243,8 +340,8 @@ public void CanAddDatabasesWithDifferentNamesOnSingleServer() Assert.Equal("customers1", db1.Resource.DatabaseName); Assert.Equal("customers2", db2.Resource.DatabaseName); - Assert.Equal("{mongo1.connectionString}/customers1", db1.Resource.ConnectionStringExpression.ValueExpression); - Assert.Equal("{mongo1.connectionString}/customers2", db2.Resource.ConnectionStringExpression.ValueExpression); + Assert.Equal("mongodb://{mongo1.bindings.tcp.host}:{mongo1.bindings.tcp.port}/customers1", db1.Resource.ConnectionStringExpression.ValueExpression); + Assert.Equal("mongodb://{mongo1.bindings.tcp.host}:{mongo1.bindings.tcp.port}/customers2", db2.Resource.ConnectionStringExpression.ValueExpression); } [Fact] @@ -261,7 +358,7 @@ public void CanAddDatabasesWithTheSameNameOnMultipleServers() Assert.Equal("imports", db1.Resource.DatabaseName); Assert.Equal("imports", db2.Resource.DatabaseName); - Assert.Equal("{mongo1.connectionString}/imports", db1.Resource.ConnectionStringExpression.ValueExpression); - Assert.Equal("{mongo2.connectionString}/imports", db2.Resource.ConnectionStringExpression.ValueExpression); + Assert.Equal("mongodb://{mongo1.bindings.tcp.host}:{mongo1.bindings.tcp.port}/imports", db1.Resource.ConnectionStringExpression.ValueExpression); + Assert.Equal("mongodb://{mongo2.bindings.tcp.host}:{mongo2.bindings.tcp.port}/imports", db2.Resource.ConnectionStringExpression.ValueExpression); } } diff --git a/tests/Aspire.Hosting.MongoDB.Tests/MongoDbFunctionalTests.cs b/tests/Aspire.Hosting.MongoDB.Tests/MongoDbFunctionalTests.cs index cdf596fb10..2887d80d7c 100644 --- a/tests/Aspire.Hosting.MongoDB.Tests/MongoDbFunctionalTests.cs +++ b/tests/Aspire.Hosting.MongoDB.Tests/MongoDbFunctionalTests.cs @@ -2,17 +2,17 @@ // The .NET Foundation licenses this file to you under the MIT license. using Aspire.Components.Common.Tests; +using Aspire.Hosting.ApplicationModel; using Aspire.Hosting.Utils; using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Diagnostics.HealthChecks; using Microsoft.Extensions.Hosting; -using MongoDB.Bson.Serialization.Attributes; using MongoDB.Bson; +using MongoDB.Bson.Serialization.Attributes; using MongoDB.Driver; +using Polly; using Xunit; using Xunit.Abstractions; -using Polly; -using Aspire.Hosting.ApplicationModel; -using Microsoft.Extensions.Diagnostics.HealthChecks; namespace Aspire.Hosting.MongoDB.Tests; @@ -102,6 +102,69 @@ await pipeline.ExecuteAsync(async token => }, cts.Token); } + [InlineData(null)] + [InlineData(10003)] + [Theory] + [RequiresDocker] + public async Task VerifyMongoDBResourceReplicaSet(int? customPort) + { + var cts = new CancellationTokenSource(TimeSpan.FromMinutes(3)); + var pipeline = new ResiliencePipelineBuilder() + .AddRetry(new() { MaxRetryAttempts = 10, Delay = TimeSpan.FromSeconds(1) }) + .Build(); + + using var builder = TestDistributedApplicationBuilder.CreateWithTestContainerRegistry(testOutputHelper); + + var mongodb = builder.AddMongoDB("mongodb", customPort) + .WithReplicaSet(); + var db = mongodb.AddDatabase("testdb"); + using var app = builder.Build(); + + await app.StartAsync(); + + var hb = Host.CreateApplicationBuilder(); + + hb.Configuration[$"ConnectionStrings:{db.Resource.Name}"] = await db.Resource.ConnectionStringExpression.GetValueAsync(default); + + hb.AddMongoDBClient(db.Resource.Name); + + using var host = hb.Build(); + + await host.StartAsync(); + + await pipeline.ExecuteAsync(async token => + { + var mongoDatabase = host.Services.GetRequiredService(); + + var collection = mongoDatabase.GetCollection(CollectionName); + var pipeline = new EmptyPipelineDefinition>() + .Match(x => x.OperationType == ChangeStreamOperationType.Insert); + + var changeStreamOptions = new ChangeStreamOptions + { + FullDocument = ChangeStreamFullDocumentOption.UpdateLookup + }; + + // This API requires replica sets + using var cursor = await collection.WatchAsync(pipeline, changeStreamOptions, cts.Token); + var tcs = new TaskCompletionSource(); + + var name = s_movies[0].Name + "Updated"; + _ = cursor.ForEachAsync(cursor => + { + Assert.Equal(ChangeStreamOperationType.Insert, cursor.OperationType); + Assert.NotNull(cursor.FullDocument); + tcs.SetResult(cursor.FullDocument); + }, cts.Token); + + await collection.InsertOneAsync(new Movie { Name = name }, cancellationToken: token); + + var updated = await tcs.Task; + + Assert.Equal(name, updated.Name); + }, cts.Token); + } + [Theory] [InlineData(true)] [InlineData(false)] diff --git a/tests/Aspire.Hosting.Tests/WaitForTests.cs b/tests/Aspire.Hosting.Tests/WaitForTests.cs index b2d4cf2370..638f380489 100644 --- a/tests/Aspire.Hosting.Tests/WaitForTests.cs +++ b/tests/Aspire.Hosting.Tests/WaitForTests.cs @@ -4,6 +4,7 @@ using Aspire.Components.Common.Tests; using Aspire.Hosting.Utils; using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Diagnostics.HealthChecks; using Xunit; using Xunit.Abstractions; @@ -153,7 +154,7 @@ await rns.PublishUpdateAsync(dependency.Resource, s => s with public async Task EnsureDependencyResourceThatReturnsNonMatchingExitCodeResultsInDependentResourceFailingToStart() { using var builder = TestDistributedApplicationBuilder.Create().WithTestAndResourceLogging(testOutputHelper); - + var dependency = builder.AddResource(new CustomResource("test")); var nginx = builder.AddContainer("nginx", "mcr.microsoft.com/cbl-mariner/base/nginx", "1.22") .WithReference(dependency) @@ -272,8 +273,78 @@ await rns.PublishUpdateAsync(dependency.Resource, s => s with await app.StopAsync(); } + [InlineData(true)] + [InlineData(false)] + [Theory] + [RequiresDocker] + public async Task WaitForWaitsForHealthChecks(bool includeHealthChecks) + { + using var builder = TestDistributedApplicationBuilder.Create().WithTestAndResourceLogging(testOutputHelper); + + var healthService = new TaskCompletedHealthCheck("test-health"); + builder.Services.AddHealthChecks() + .Add(new HealthCheckRegistration(healthService.Name, healthService, default, default)); + + var dependency = builder.AddResource(new CustomResource("test")) + .WithHealthCheck(healthService.Name); + + var nginx = builder.AddContainer("nginx", "mcr.microsoft.com/cbl-mariner/base/nginx", "1.22") + .WithReference(dependency) + .WaitFor(dependency, includeHealthChecks: includeHealthChecks); + + using var app = builder.Build(); + + // StartAsync will currently block until the dependency resource moves + // into a Finished state, so rather than awaiting it we'll hold onto the + // task so we can inspect the state of the Nginx resource which should + // be in a waiting state if everything is working correctly. + var startupCts = new CancellationTokenSource(TimeSpan.FromSeconds(60)); + var startTask = app.StartAsync(startupCts.Token); + + // We don't want to wait forever for Nginx to move into a waiting state, + // it should be super quick, but we'll allow 60 seconds just in case the + // CI machine is chugging (also useful when collecting code coverage). + var waitingStateCts = new CancellationTokenSource(TimeSpan.FromSeconds(60)); + + var rns = app.Services.GetRequiredService(); + await rns.WaitForResourceAsync(nginx.Resource.Name, KnownResourceStates.Waiting, waitingStateCts.Token); + + // Now that we know we successfully entered the Waiting state, we can start the dependency + await rns.PublishUpdateAsync(dependency.Resource, s => s with + { + State = KnownResourceStates.Running + }); + + if (includeHealthChecks) + { + healthService.SetResult(); + } + + await rns.WaitForResourceAsync(nginx.Resource.Name, KnownResourceStates.Running, waitingStateCts.Token); + + await startTask; + + await app.StopAsync(); + } + private sealed class CustomResource(string name) : Resource(name), IResourceWithConnectionString { public ReferenceExpression ConnectionStringExpression => ReferenceExpression.Create($"foo"); } + + private sealed class TaskCompletedHealthCheck(string name) : TaskCompletionSource, IHealthCheck + { + public string Name => name; + + public Task CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default) => Task.FromResult(CheckHealth()); + + private HealthCheckResult CheckHealth() => Task switch + { + { IsCompletedSuccessfully: true } => HealthCheckResult.Healthy(), + { IsFaulted: true } => HealthCheckResult.Unhealthy("Error", Task.Exception), + { IsCanceled: true } => HealthCheckResult.Unhealthy("Canceled"), + { IsCompleted: true } => HealthCheckResult.Unhealthy("Still running"), + _ => HealthCheckResult.Unhealthy("Unknown") + }; + } }