From b8caefdc6eac62772d29681d8f9fee437fdfe9c3 Mon Sep 17 00:00:00 2001
From: Julie Lycklama <8868475+jnlycklama@users.noreply.github.com>
Date: Mon, 30 Oct 2023 11:45:47 -0700
Subject: [PATCH] Use abstract AzureStorageHealthCheck to test for CMK health
 (#3155)

* upgrade shared components

* revert bad change from merge
---
 .../HealthChecks/DurableTaskHealthCheck.cs    | 23 +++----------------
 1 file changed, 3 insertions(+), 20 deletions(-)

diff --git a/src/Microsoft.Health.Dicom.Functions.Client/HealthChecks/DurableTaskHealthCheck.cs b/src/Microsoft.Health.Dicom.Functions.Client/HealthChecks/DurableTaskHealthCheck.cs
index 078779f44a..fdfc2b285a 100644
--- a/src/Microsoft.Health.Dicom.Functions.Client/HealthChecks/DurableTaskHealthCheck.cs
+++ b/src/Microsoft.Health.Dicom.Functions.Client/HealthChecks/DurableTaskHealthCheck.cs
@@ -3,7 +3,6 @@
 // Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
 // -------------------------------------------------------------------------------------------------
 
-using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 using EnsureThat;
@@ -15,39 +14,23 @@
 
 namespace Microsoft.Health.Dicom.Functions.Client.HealthChecks;
 
-internal sealed class DurableTaskHealthCheck : IHealthCheck
+internal sealed class DurableTaskHealthCheck : AzureStorageHealthCheck
 {
-    private const string DegradedDescription = "The health of the task hub has degraded.";
-
     private readonly ITaskHubClient _client;
-    private readonly ValueCache<CustomerKeyHealth> _customerKeyHealthCache;
     private readonly ILogger<DurableTaskHealthCheck> _logger;
 
     public DurableTaskHealthCheck(
         ITaskHubClient client,
         ValueCache<CustomerKeyHealth> customerKeyHealthCache,
         ILogger<DurableTaskHealthCheck> logger)
+        : base(customerKeyHealthCache, logger)
     {
         _client = EnsureArg.IsNotNull(client, nameof(client));
-        _customerKeyHealthCache = EnsureArg.IsNotNull(customerKeyHealthCache, nameof(customerKeyHealthCache));
         _logger = EnsureArg.IsNotNull(logger, nameof(logger));
     }
 
-    public async Task<HealthCheckResult> CheckHealthAsync(HealthCheckContext context, CancellationToken cancellationToken = default)
+    public override async Task<HealthCheckResult> CheckAzureStorageHealthAsync(CancellationToken cancellationToken)
     {
-        EnsureArg.IsNotNull(context, nameof(context));
-
-        CustomerKeyHealth cmkStatus = await _customerKeyHealthCache.GetAsync(cancellationToken).ConfigureAwait(false);
-        if (!cmkStatus.IsHealthy)
-        {
-            // if the customer-managed key is inaccessible, the task hub will also be inaccessible
-            return new HealthCheckResult(
-                HealthStatus.Degraded,
-                DegradedDescription,
-                cmkStatus.Exception,
-                new Dictionary<string, object> { { "Reason", cmkStatus.Reason } });
-        }
-
         ITaskHub taskHub = await _client.GetTaskHubAsync(cancellationToken).ConfigureAwait(false);
         if (taskHub == null)
             return HealthCheckResult.Unhealthy("Task hub not found.");