Add WireMockHealthCheck in WireMock.Net.Aspire

examples-Aspire/AspireApp1.AppHost/Program.cs

@@ -45,6 +45,7 @@
 
 builder.AddProject<Projects.AspireApp1_Web>("webfrontend")
     .WithExternalHttpEndpoints()
-    .WithReference(apiService);
+    .WithReference(apiService)
+    .WaitFor(apiService);
 
 builder.Build().Run();

src/WireMock.Net.Aspire/WireMockHealthCheck.cs

@@ -0,0 +1,44 @@
+// Copyright © WireMock.Net
+
+using Aspire.Hosting.ApplicationModel;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+using WireMock.Client;
+
+namespace WireMock.Net.Aspire;
+
+/// <summary>
+/// WireMockHealthCheck
+/// </summary>
+public class WireMockHealthCheck(WireMockServerResource resource) : IHealthCheck
+{
+    private const string HealthStatusHealthy = "Healthy";
+
+    /// <inheritdoc />
+    public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
+    {
+        if (!await IsHealthyAsync(resource.AdminApi.Value, cancellationToken))
+        {
+            return HealthCheckResult.Unhealthy("WireMock.Net is not healthy");
+        }
+
+        if (resource.ApiMappingState == WireMockMappingState.NotSubmitted)
+        {
+            return HealthCheckResult.Unhealthy("WireMock.Net has not received mappings");
+        }
+
+        return HealthCheckResult.Healthy();
+    }
+
+    private static async Task<bool> IsHealthyAsync(IWireMockAdminApi adminApi, CancellationToken cancellationToken)
+    {
+        try
+        {
+            var status = await adminApi.GetHealthAsync(cancellationToken);
+            return string.Equals(status, HealthStatusHealthy, StringComparison.OrdinalIgnoreCase);
+        }
+        catch
+        {
+            return false;
+        }
+    }
+}

src/WireMock.Net.Aspire/WireMockMappingState.cs

@@ -0,0 +1,10 @@
+// Copyright © WireMock.Net
+
+namespace WireMock.Net.Aspire;
+
+internal enum WireMockMappingState
+{
+    NoMappings,
+    NotSubmitted,
+    Submitted,
+}

src/WireMock.Net.Aspire/WireMockServerBuilderExtensions.cs

@@ -3,6 +3,7 @@
 using Aspire.Hosting.ApplicationModel;
 using Aspire.Hosting.Lifecycle;
 using Aspire.Hosting.WireMock;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Diagnostics.HealthChecks;
 using Stef.Validation;
 using WireMock.Client.Builders;
@@ -53,11 +54,21 @@ public static IResourceBuilder<WireMockServerResource> AddWireMock(this IDistrib
         Guard.NotNull(arguments);
 
         var wireMockContainerResource = new WireMockServerResource(name, arguments);
+
+        var healthCheckKey = $"{name}_check";
+        var healthCheckRegistration = new HealthCheckRegistration(
+            healthCheckKey,
+            _ => new WireMockHealthCheck(wireMockContainerResource),
+            failureStatus: null,
+            tags: null);
+        builder.Services.AddHealthChecks().Add(healthCheckRegistration);
+
         var resourceBuilder = builder
             .AddResource(wireMockContainerResource)
             .WithImage(DefaultLinuxImage)
             .WithEnvironment(ctx => ctx.EnvironmentVariables.Add("DOTNET_USE_POLLING_FILE_WATCHER", "1")) // https://khalidabuhakmeh.com/aspnet-docker-gotchas-and-workarounds#configuration-reloads-and-filesystemwatcher
             .WithHttpEndpoint(port: arguments.HttpPort, targetPort: WireMockServerArguments.HttpContainerPort)
+            .WithHealthCheck(healthCheckKey)
             .WithWireMockInspectorCommand();
 
         if (!string.IsNullOrEmpty(arguments.MappingsPath))
@@ -172,6 +183,7 @@ public static IResourceBuilder<WireMockServerResource> WithApiMappingBuilder(thi
 
         wiremock.ApplicationBuilder.Services.TryAddLifecycleHook<WireMockServerLifecycleHook>();
         wiremock.Resource.Arguments.ApiMappingBuilder = configure;
+        wiremock.Resource.ApiMappingState = WireMockMappingState.NotSubmitted;
 
         return wiremock;
     }

src/WireMock.Net.Aspire/WireMockServerLifecycleHook.cs

@@ -10,32 +10,47 @@ internal class WireMockServerLifecycleHook(ILoggerFactory loggerFactory) : IDist
 {
     private readonly CancellationTokenSource _shutdownCts = new();
 
-    public async Task AfterResourcesCreatedAsync(DistributedApplicationModel appModel, CancellationToken cancellationToken = default)
-    {
-        var cts = CancellationTokenSource.CreateLinkedTokenSource(_shutdownCts.Token, cancellationToken);
+    private CancellationTokenSource? _linkedCts;
+    private Task? _mappingTask;
 
-        var wireMockServerResources = appModel.Resources
-            .OfType<WireMockServerResource>()
-            .ToArray();
+    public Task AfterEndpointsAllocatedAsync(DistributedApplicationModel appModel, CancellationToken cancellationToken = default)
+    {
+        _linkedCts = CancellationTokenSource.CreateLinkedTokenSource(_shutdownCts.Token, cancellationToken);
 
-        foreach (var wireMockServerResource in wireMockServerResources)
+        _mappingTask = Task.Run(async () =>
         {
-            wireMockServerResource.SetLogger(loggerFactory.CreateLogger<WireMockServerResource>());
+            var wireMockServerResources = appModel.Resources
+                .OfType<WireMockServerResource>()
+                .ToArray();
 
-            var endpoint = wireMockServerResource.GetEndpoint();
-            if (endpoint.IsAllocated)
+            foreach (var wireMockServerResource in wireMockServerResources)
             {
-                await wireMockServerResource.WaitForHealthAsync(cts.Token);
+                wireMockServerResource.SetLogger(loggerFactory.CreateLogger<WireMockServerResource>());
+
+                var endpoint = wireMockServerResource.GetEndpoint();
+                System.Diagnostics.Debug.Assert(endpoint.IsAllocated);
 
-                await wireMockServerResource.CallApiMappingBuilderActionAsync(cts.Token);
+                await wireMockServerResource.WaitForHealthAsync(_linkedCts.Token);
 
-                wireMockServerResource.StartWatchingStaticMappings(cts.Token);
+                await wireMockServerResource.CallApiMappingBuilderActionAsync(_linkedCts.Token);
+
+                wireMockServerResource.StartWatchingStaticMappings(_linkedCts.Token);
             }
-        }
+        }, _linkedCts.Token);
+
+        return Task.CompletedTask;
     }
 
     public async ValueTask DisposeAsync()
     {
         await _shutdownCts.CancelAsync();
+
+        _linkedCts?.Dispose();
+        _shutdownCts.Dispose();
+
+        if (_mappingTask is not null)
+        {
+            await _mappingTask;
+        }
     }
 }

src/WireMock.Net.Aspire/WireMockServerResource.cs

@@ -5,6 +5,7 @@
 using Stef.Validation;
 using WireMock.Client;
 using WireMock.Client.Extensions;
+using WireMock.Net.Aspire;
 using WireMock.Util;
 
 // ReSharper disable once CheckNamespace
@@ -19,6 +20,7 @@ public class WireMockServerResource : ContainerResource, IResourceWithServiceDis
 
     internal WireMockServerArguments Arguments { get; }
     internal Lazy<IWireMockAdminApi> AdminApi => new(CreateWireMockAdminApi);
+    internal WireMockMappingState ApiMappingState { get; set; } = WireMockMappingState.NoMappings;
 
     private ILogger? _logger;
     private EnhancedFileSystemWatcher? _enhancedFileSystemWatcher;
@@ -64,6 +66,8 @@ internal async Task CallApiMappingBuilderActionAsync(CancellationToken cancellat
 
         var mappingBuilder = AdminApi.Value.GetMappingBuilder();
         await Arguments.ApiMappingBuilder.Invoke(mappingBuilder, cancellationToken);
+
+        ApiMappingState = WireMockMappingState.Submitted;
     }
 
     internal void StartWatchingStaticMappings(CancellationToken cancellationToken)

test/WireMock.Net.Aspire.Tests/IntegrationTests.cs

@@ -8,7 +8,7 @@
 
 namespace WireMock.Net.Aspire.Tests;
 
 public class IntegrationTests(ITestOutputHelper output)
 {
     private record WeatherForecast(DateOnly Date, int TemperatureC, string? Summary);
 
@@ -19,6 +19,7 @@
         var appHostBuilder = await DistributedApplicationTestingBuilder.CreateAsync<WireMock_Net_Aspire_TestAppHost>();
         await using var app = await appHostBuilder.BuildAsync();
         await app.StartAsync();
+        await app.ResourceNotifications.WaitForResourceHealthyAsync("wiremock-service");
 
         using var httpClient = app.CreateHttpClient("wiremock-service");
 
@@ -46,6 +47,7 @@
         var appHostBuilder = await DistributedApplicationTestingBuilder.CreateAsync<WireMock_Net_Aspire_TestAppHost>();
         await using var app = await appHostBuilder.BuildAsync();
         await app.StartAsync();
+        await app.ResourceNotifications.WaitForResourceHealthyAsync("wiremock-service");
 
         var adminClient = app.CreateWireMockAdminClient("wiremock-service");
 

test/WireMock.Net.Aspire.Tests/WireMockServerBuilderExtensionsTests.cs

@@ -67,7 +67,7 @@ public void AddWireMock()
             MappingsPath = null,
             HttpPort = port
         });
-        wiremock.Resource.Annotations.Should().HaveCount(5);
+        wiremock.Resource.Annotations.Should().HaveCount(6);
 
         var containerImageAnnotation = wiremock.Resource.Annotations.OfType<ContainerImageAnnotation>().FirstOrDefault();
         containerImageAnnotation.Should().BeEquivalentTo(new ContainerImageAnnotation