Added managed identity for Azure IoTHub health checks

src/HealthChecks.Azure.IoTHub/DependencyInjection/IoTHubHealthChecksBuilderExtensions.cs

@@ -1,19 +1,25 @@
 using HealthChecks.Azure.IoTHub;
+using Microsoft.Azure.Devices;
 using Microsoft.Extensions.Diagnostics.HealthChecks;
 
 namespace Microsoft.Extensions.DependencyInjection;
 
 /// <summary>
-/// Extension methods to configure <see cref="IoTHubHealthCheck"/>.
+/// Extension methods to configure <see cref="IoTHubRegistryManagerHealthCheck"/>.
 /// </summary>
 public static class IoTHubHealthChecksBuilderExtensions
 {
-    private const string NAME = "iothub";
+    private const string NAME_REGISTRY_MANAGER = "iothub_registrymanager";
+    private const string NAME_SERVICE_CLIENT = "iothub_serviceclient";
 
     /// <summary>
-    /// Add a health check for Azure IoT Hub.
+    /// Add a health check for Azure IoT Hub registry manager.
     /// </summary>
     /// <param name="builder">The <see cref="IHealthChecksBuilder"/>.</param>
+    /// <param name="registryManagerFactory">
+    /// /// An optional factory to obtain <see cref="RegistryManager" /> instance.
+    /// When not provided, <see cref="RegistryManager" /> is simply resolved from <see cref="IServiceProvider"/>.
+    /// </param>
     /// <param name="optionsFactory">A action to configure the Azure IoT Hub connection to use.</param>
     /// <param name="name">The health check name. Optional. If <c>null</c> the type name 'iothub' will be used for the name.</param>
     /// <param name="failureStatus">
@@ -23,22 +29,56 @@ public static class IoTHubHealthChecksBuilderExtensions
     /// <param name="tags">A list of tags that can be used to filter sets of health checks. Optional.</param>
     /// <param name="timeout">An optional <see cref="TimeSpan"/> representing the timeout of the check.</param>
     /// <returns>The specified <paramref name="builder"/>.</returns>
-    public static IHealthChecksBuilder AddAzureIoTHub(
+    public static IHealthChecksBuilder AddAzureIoTHubRegistryManager(
         this IHealthChecksBuilder builder,
-        Action<IoTHubOptions>? optionsFactory,
+        Func<IServiceProvider, RegistryManager>? registryManagerFactory = default,
+        Func<IServiceProvider, IotHubRegistryManagerOptions>? optionsFactory = default,
         string? name = default,
         HealthStatus? failureStatus = default,
         IEnumerable<string>? tags = default,
         TimeSpan? timeout = default)
     {
-        var options = new IoTHubOptions();
-        optionsFactory?.Invoke(options);
-
-        var registrationName = name ?? NAME;
+        return builder.Add(new HealthCheckRegistration(
+           name ?? NAME_REGISTRY_MANAGER,
+           sp => new IoTHubRegistryManagerHealthCheck(
+                registryManager: registryManagerFactory?.Invoke(sp) ?? sp.GetRequiredService<RegistryManager>(),
+                options: optionsFactory?.Invoke(sp)),
+           failureStatus,
+           tags,
+           timeout));
+    }
 
+    /// <summary>
+    /// Add a health check for Azure IoT Hub service client.
+    /// </summary>
+    /// <param name="builder">The <see cref="IHealthChecksBuilder"/>.</param>
+    /// <param name="serviceClientFactory">
+    /// /// An optional factory to obtain <see cref="ServiceClient" /> instance.
+    /// When not provided, <see cref="ServiceClient" /> is simply resolved from <see cref="IServiceProvider"/>.
+    /// </param>
+    /// <param name="optionsFactory">A action to configure the Azure IoT Hub connection to use.</param>
+    /// <param name="name">The health check name. Optional. If <c>null</c> the type name 'iothub' will be used for the name.</param>
+    /// <param name="failureStatus">
+    /// The <see cref="HealthStatus"/> that should be reported when the health check fails. Optional. If <c>null</c> then
+    /// the default status of <see cref="HealthStatus.Unhealthy"/> will be reported.
+    /// </param>
+    /// <param name="tags">A list of tags that can be used to filter sets of health checks. Optional.</param>
+    /// <param name="timeout">An optional <see cref="TimeSpan"/> representing the timeout of the check.</param>
+    /// <returns>The specified <paramref name="builder"/>.</returns>
+    public static IHealthChecksBuilder AddAzureIoTHubServiceClient(
+        this IHealthChecksBuilder builder,
+        Func<IServiceProvider, ServiceClient>? serviceClientFactory = default,
+        Func<IServiceProvider, IotHubServiceClientOptions>? optionsFactory = default,
+        string? name = default,
+        HealthStatus? failureStatus = default,
+        IEnumerable<string>? tags = default,
+        TimeSpan? timeout = default)
+    {
         return builder.Add(new HealthCheckRegistration(
-           registrationName,
-           sp => new IoTHubHealthCheck(options),
+           name ?? NAME_SERVICE_CLIENT,
+           sp => new IoTHubServiceClientHealthCheck(
+                serviceClient: serviceClientFactory?.Invoke(sp) ?? sp.GetRequiredService<ServiceClient>(),
+                options: optionsFactory?.Invoke(sp)),
            failureStatus,
            tags,
            timeout));

src/HealthChecks.Azure.IoTHub/IoTHubHealthCheck.cs → src/HealthChecks.Azure.IoTHub/IoTHubRegistryManagerHealthCheck.cs

@@ -3,18 +3,25 @@
 
 namespace HealthChecks.Azure.IoTHub;
 
-public class IoTHubHealthCheck : IHealthCheck
+public sealed class IoTHubRegistryManagerHealthCheck : IHealthCheck
 {
-    private readonly IoTHubOptions _options;
+    private readonly IotHubRegistryManagerOptions _options;
+    private readonly RegistryManager _registryManager;
 
-    public IoTHubHealthCheck(IoTHubOptions options)
+    public IoTHubRegistryManagerHealthCheck(RegistryManager registryManager, IotHubRegistryManagerOptions? options = default)
     {
-        _options = Guard.ThrowIfNull(options);
+        _options = options ?? new IotHubRegistryManagerOptions();
+        _registryManager = Guard.ThrowIfNull(registryManager);
     }
 
     /// <inheritdoc />
     public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
     {
+        if (!_options.RegistryWriteCheck && !_options.RegistryReadCheck)
+        {
+            return new HealthCheckResult(context.Registration.FailureStatus, description: $"No health check enabled, both {nameof(IotHubRegistryManagerOptions.RegistryReadCheck)} and {nameof(IotHubRegistryManagerOptions.RegistryWriteCheck)} are false");
+        }
+
         try
         {
             if (_options.RegistryWriteCheck)
@@ -25,10 +32,6 @@ public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context
             {
                 await ExecuteRegistryReadCheckAsync().ConfigureAwait(false);
             }
-            if (_options.ServiceConnectionCheck)
-            {
-                await ExecuteServiceConnectionCheckAsync(cancellationToken).ConfigureAwait(false);
-            }
 
             return HealthCheckResult.Healthy();
         }
@@ -38,37 +41,29 @@ public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context
         }
     }
 
-    private async Task ExecuteServiceConnectionCheckAsync(CancellationToken cancellationToken)
-    {
-        using var client = ServiceClient.CreateFromConnectionString(_options.ConnectionString, _options.ServiceConnectionTransport);
-        await client.GetServiceStatisticsAsync(cancellationToken).ConfigureAwait(false);
-    }
 
     private async Task ExecuteRegistryReadCheckAsync()
     {
-        using var client = RegistryManager.CreateFromConnectionString(_options.ConnectionString);
-        var query = client.CreateQuery(_options.RegistryReadQuery, 1);
+        var query = _registryManager.CreateQuery(_options.RegistryReadQuery, 1);
         await query.GetNextAsJsonAsync().ConfigureAwait(false);
     }
 
     private async Task ExecuteRegistryWriteCheckAsync(CancellationToken cancellationToken)
     {
-        using var client = RegistryManager.CreateFromConnectionString(_options.ConnectionString);
-
         var deviceId = _options.RegistryWriteDeviceIdFactory();
-        var device = await client.GetDeviceAsync(deviceId, cancellationToken).ConfigureAwait(false);
+        var device = await _registryManager.GetDeviceAsync(deviceId, cancellationToken).ConfigureAwait(false);
 
         // in default implementation of configuration deviceId equals "health-check-registry-write-device-id"
         // if in previous health check device were not removed -- try remove it
         // if in previous health check device were added and removed -- try create and remove it
         if (device != null)
         {
-            await client.RemoveDeviceAsync(deviceId, cancellationToken).ConfigureAwait(false);
+            await _registryManager.RemoveDeviceAsync(deviceId, cancellationToken).ConfigureAwait(false);
         }
         else
         {
-            await client.AddDeviceAsync(new Device(deviceId), cancellationToken).ConfigureAwait(false);
-            await client.RemoveDeviceAsync(deviceId, cancellationToken).ConfigureAwait(false);
+            await _registryManager.AddDeviceAsync(new Device(deviceId), cancellationToken).ConfigureAwait(false);
+            await _registryManager.RemoveDeviceAsync(deviceId, cancellationToken).ConfigureAwait(false);
         }
     }
 }

src/HealthChecks.Azure.IoTHub/IoTHubServiceClientHealthCheck.cs

@@ -0,0 +1,36 @@
+using Microsoft.Azure.Devices;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+
+namespace HealthChecks.Azure.IoTHub;
+
+public sealed class IoTHubServiceClientHealthCheck : IHealthCheck
+{
+    private readonly IotHubServiceClientOptions _options;
+    private readonly ServiceClient _serviceClient;
+
+    public IoTHubServiceClientHealthCheck(ServiceClient serviceClient, IotHubServiceClientOptions? options = default)
+    {
+        _options = options ?? new IotHubServiceClientOptions();
+        _serviceClient = Guard.ThrowIfNull(serviceClient);
+    }
+
+    /// <inheritdoc />
+    public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
+    {
+        try
+        {
+            await ExecuteServiceConnectionCheckAsync(cancellationToken).ConfigureAwait(false);
+
+            return HealthCheckResult.Healthy();
+        }
+        catch (Exception ex)
+        {
+            return new HealthCheckResult(context.Registration.FailureStatus, exception: ex);
+        }
+    }
+
+    private async Task ExecuteServiceConnectionCheckAsync(CancellationToken cancellationToken)
+    {
+        await _serviceClient.GetServiceStatisticsAsync(cancellationToken).ConfigureAwait(false);
+    }
+}

src/HealthChecks.Azure.IoTHub/IotHubRegistryManagerOptions.cs

@@ -0,0 +1,23 @@
+namespace HealthChecks.Azure.IoTHub;
+
+public sealed class IotHubRegistryManagerOptions
+{
+    internal bool RegistryReadCheck { get; private set; }
+    internal bool RegistryWriteCheck { get; private set; }
+    internal string RegistryReadQuery { get; private set; } = null!;
+    internal Func<string> RegistryWriteDeviceIdFactory { get; private set; } = null!;
+
+    public IotHubRegistryManagerOptions AddRegistryReadCheck(string query = "SELECT deviceId FROM devices")
+    {
+        RegistryReadCheck = true;
+        RegistryReadQuery = query;
+        return this;
+    }
+
+    public IotHubRegistryManagerOptions AddRegistryWriteCheck(Func<string>? deviceIdFactory = null)
+    {
+        RegistryWriteCheck = true;
+        RegistryWriteDeviceIdFactory = deviceIdFactory ?? (() => "health-check-registry-write-device-id");
+        return this;
+    }
+}

src/HealthChecks.Azure.IoTHub/IotHubServiceClientOptions.cs

@@ -0,0 +1,2 @@
+
+public sealed class IotHubServiceClientOptions;

src/HealthChecks.Azure.IoTHub/README.md

@@ -1,27 +1,44 @@
-# Azure IoT Hub Health Check
+## Azure IoT Hub Health Check
 
 This health check verifies the ability to communicate with Azure IoT Hub. For more information about Azure IoT Hub please check and .NET please check the [Azure IoT Hub Microsoft Site](https://azure.microsoft.com/en-us/services/iot-hub/)
 
-## Example Usage
-
-With all of the following examples, you can additionally add the following parameters:
-
-- `name`: The health check name. Default if not specified is `iothub`.
-- `failureStatus`: The `HealthStatus` that should be reported when the health check fails. Default is `HealthStatus.Unhealthy`.
-- `tags`: A list of tags that can be used to filter sets of health checks.
-- `timeout`: A `System.TimeSpan` representing the timeout of the check.
-
-### Basic
+### Defaults
+note: AddRegistryReadCheck() or AddRegistryWriteCheck() is required to be called else it will return unhealthy for the Registry manager health check
 
 ```csharp
 public void ConfigureServices(IServiceCollection services)
 {
+    Services.AddSingleton(sp => RegistryManager.Create("iot-hub-hostname", new DefaultAzureCredential());
+    Services.AddSingleton(sp => ServiceClient.Create("iot-hub-hostname", new DefaultAzureCredential());
+
     services
         .AddHealthChecks()
-        .AddAzureIoTHub(options =>
-        {
-            options.AddConnectionString("iot-hub-connectionstring")
-                .AddServiceConnectionCheck();
-        });
+        .AddAzureIoTHubRegistryManager(
+            clientFactory: sp.GetRequiredService<RegistryManager>()
+            optionsFactory: sp => new IotHubRegistryManagerOptions()
+              .AddRegistryReadCheck()
+              .AddRegistryWriteCheck();
+       .AddAzureIoTHubServiceClient();
 }
 ```
+
+
+### Customization
+
+With all of the following examples, you can additionally add the following parameters:
+
+AddAzureIoTHubServiceClient
+- `serviceClientFactory`: A factory method to provide `ServiceClient` instance.
+- `optionsFactory`: A factory method to provide `IotHubServiceClientOptions` instance. It allows to specify the secret name and whether the secret should be created when it's not found.
+- `name`: The health check name. Default if not specified is `iothub`.
+- `failureStatus`: The `HealthStatus` that should be reported when the health check fails. Default is `HealthStatus.Unhealthy`.
+- `tags`: A list of tags that can be used to filter sets of health checks.
+- `timeout`: A `System.TimeSpan` representing the timeout of the check.
+
+AddAzureIoTHubRegistryManager
+- `registryManagerFactory`: A factory method to provide `RegistryManager` instance.
+- `optionsFactory`: A factory method to provide `IotHubRegistryManagerOptions` instance. It allows to specify the secret name and whether the secret should be created when it's not found.
+- `name`: The health check name. Default if not specified is `iothub`.
+- `failureStatus`: The `HealthStatus` that should be reported when the health check fails. Default is `HealthStatus.Unhealthy`.
+- `tags`: A list of tags that can be used to filter sets of health checks.
+- `timeout`: A `System.TimeSpan` representing the timeout of the check.

test/HealthChecks.Azure.IoTHub.Tests/HealthChecks.Azure.IoTHub.Tests.csproj

@@ -1,6 +1,7 @@
 <Project Sdk="Microsoft.NET.Sdk">
 
   <ItemGroup>
+    <PackageReference Include="Azure.Identity" Version="1.10.5" />
     <ProjectReference Include="..\..\src\HealthChecks.Azure.IoTHub\HealthChecks.Azure.IoTHub.csproj" />
   </ItemGroup>