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

examples-Aspire/AspireApp1.AppHost/Program.cs

@@ -45,6 +45,7 @@ IResourceBuilder<WireMockServerResource> apiService = builder
 
 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 class WireMockServerBuilderExtensions
         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 class WireMockServerBuilderExtensions
 
         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)
+    private CancellationTokenSource? _linkedCts;
+    private Task? _mappingTask;
+
+    public Task AfterEndpointsAllocatedAsync(DistributedApplicationModel appModel, CancellationToken cancellationToken = default)
     {
-        var cts = CancellationTokenSource.CreateLinkedTokenSource(_shutdownCts.Token, cancellationToken);
+        _linkedCts = CancellationTokenSource.CreateLinkedTokenSource(_shutdownCts.Token, cancellationToken);
 
-        var wireMockServerResources = appModel.Resources
+        _mappingTask = Task.Run(async () =>
-            .OfType<WireMockServerResource>()
-            .ToArray();
-
-        foreach (var wireMockServerResource in wireMockServerResources)
         {
-            wireMockServerResource.SetLogger(loggerFactory.CreateLogger<WireMockServerResource>());
+            var wireMockServerResources = appModel.Resources
+                .OfType<WireMockServerResource>()
+                .ToArray();
 
-            var endpoint = wireMockServerResource.GetEndpoint();
+            foreach (var wireMockServerResource in wireMockServerResources)
-            if (endpoint.IsAllocated)
             {
-                await wireMockServerResource.WaitForHealthAsync(cts.Token);
+                wireMockServerResource.SetLogger(loggerFactory.CreateLogger<WireMockServerResource>());
 
-                await wireMockServerResource.CallApiMappingBuilderActionAsync(cts.Token);
+                var endpoint = wireMockServerResource.GetEndpoint();
+                System.Diagnostics.Debug.Assert(endpoint.IsAllocated);
 
-                wireMockServerResource.StartWatchingStaticMappings(cts.Token);
+                await wireMockServerResource.WaitForHealthAsync(_linkedCts.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 RestEase;
 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 @@ public class WireMockServerResource : ContainerResource, IResourceWithServiceDis
 
         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 class IntegrationTests(ITestOutputHelper output)
         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 class IntegrationTests(ITestOutputHelper output)
         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 class WireMockServerBuilderExtensionsTests
             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