Add WireMockHealthCheck

For use with Aspire, to make WaitFor(wiremock) more useful.
Calls /__admin/health and checks the result, as well as checks if mappings using AdminApiMappingBuilder has been submitted to the server.

This created a catch-22 problem where the mappings were not submitted until the health check was healthy, but the health check was not healthy until the mappings were submitted.

To avoid this, the WireMockServerLifecycleHook class has been slightly re-arranged, and is now using the AfterEndpointsAllocatedAsync callback rather than the AfterResourcesCreatedAsync callback. Within which a separate Task is created that waits until the server is ready and submits the mappings.

Author: Zguy

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 == WireMockServerResource.MappingState.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/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 = WireMockServerResource.MappingState.NotSubmitted;
 
         return wiremock;
     }

src/WireMock.Net.Aspire/WireMockServerLifecycleHook.cs

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

src/WireMock.Net.Aspire/WireMockServerResource.cs

@@ -20,6 +20,14 @@ public class WireMockServerResource : ContainerResource, IResourceWithServiceDis
     internal WireMockServerArguments Arguments { get; }
     internal Lazy<IWireMockAdminApi> AdminApi => new(CreateWireMockAdminApi);
 
+    internal enum MappingState
+    {
+        NoMappings,
+        NotSubmitted,
+        Submitted,
+    }
+    internal MappingState ApiMappingState { get; set; } = MappingState.NoMappings;
+
     private ILogger? _logger;
     private EnhancedFileSystemWatcher? _enhancedFileSystemWatcher;
 
@@ -64,6 +72,8 @@ internal async Task CallApiMappingBuilderActionAsync(CancellationToken cancellat
 
         var mappingBuilder = AdminApi.Value.GetMappingBuilder();
         await Arguments.ApiMappingBuilder.Invoke(mappingBuilder, cancellationToken);
+
+        ApiMappingState = MappingState.Submitted;
     }
 
     internal void StartWatchingStaticMappings(CancellationToken cancellationToken)

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

@@ -19,6 +19,7 @@ public async Task StartAppHostWithWireMockAndCreateHttpClientToCallTheMockedWeat
         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 @@ public async Task StartAppHostWithWireMockAndCreateWireMockAdminClientToCallTheA
         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