Add WireMockHealthCheck in WireMock.Net.Aspire (#1375)

* 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.

* Move WireMockMappingState to its own file

* Dispose the cancellation tokens in WireMockServerLifecycleHook

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 == 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

@@ -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