From 4a9a96cef7dc055d2dd874c86a0a22d637573280 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ce=CC=81dric=20Luthi?= <cedric.luthi@gmail.com>
Date: Thu, 13 Dec 2018 17:36:01 +0100
Subject: [PATCH] Make the ResultStatusCodes property null-resettable

---
 .../HealthChecks/src/HealthCheckOptions.cs    | 45 ++++++++++---
 .../UnitTests/HealthCheckMiddlewareTests.cs   | 65 +++++++++++++++++++
 2 files changed, 102 insertions(+), 8 deletions(-)

diff --git a/src/Middleware/HealthChecks/src/HealthCheckOptions.cs b/src/Middleware/HealthChecks/src/HealthCheckOptions.cs
index d4e63cee97..c5ef255662 100644
--- a/src/Middleware/HealthChecks/src/HealthCheckOptions.cs
+++ b/src/Middleware/HealthChecks/src/HealthCheckOptions.cs
@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.Diagnostics.HealthChecks;
@@ -24,17 +25,45 @@ namespace Microsoft.AspNetCore.Diagnostics.HealthChecks
         /// </remarks>
         public Func<HealthCheckRegistration, bool> Predicate { get; set; }
 
-        /// <summary>
-        /// Gets a dictionary mapping the <see cref="HealthStatus"/> to an HTTP status code applied to the response.
-        /// This property can be used to configure the status codes returned for each status.
-        /// </summary>
-        public IDictionary<HealthStatus, int> ResultStatusCodes { get; set; } = new Dictionary<HealthStatus, int>()
+        private IDictionary<HealthStatus, int> _resultStatusCodes = new Dictionary<HealthStatus, int>(DefaultStatusCodesMapping);
+
+        private static readonly IReadOnlyDictionary<HealthStatus, int> DefaultStatusCodesMapping = new Dictionary<HealthStatus, int>
         {
-            { HealthStatus.Healthy, StatusCodes.Status200OK },
-            { HealthStatus.Degraded, StatusCodes.Status200OK },
-            { HealthStatus.Unhealthy, StatusCodes.Status503ServiceUnavailable },
+            {HealthStatus.Healthy, StatusCodes.Status200OK},
+            {HealthStatus.Degraded, StatusCodes.Status200OK},
+            {HealthStatus.Unhealthy, StatusCodes.Status503ServiceUnavailable},
         };
 
+        /// <summary>
+        /// Gets or sets a dictionary mapping the <see cref="HealthStatus"/> to an HTTP status code applied
+        /// to the response. This property can be used to configure the status codes returned for each status.
+        /// </summary>
+        /// <remarks>
+        /// Setting this property to <c>null</c> resets the mapping to its default value which maps
+        /// <see cref="HealthStatus.Healthy"/> to 200 (OK), <see cref="HealthStatus.Degraded"/> to 200 (OK) and
+        /// <see cref="HealthStatus.Unhealthy"/> to 503 (Service Unavailable).
+        /// </remarks>
+        /// <exception cref="InvalidOperationException">Thrown if at least one <see cref="HealthStatus"/> is missing when setting this property.</exception>
+        public IDictionary<HealthStatus, int> ResultStatusCodes
+        {
+            get => _resultStatusCodes;
+            set => _resultStatusCodes = value != null ? ValidateStatusCodesMapping(value) : new Dictionary<HealthStatus, int>(DefaultStatusCodesMapping);
+        }
+
+        private static IDictionary<HealthStatus, int> ValidateStatusCodesMapping(IDictionary<HealthStatus,int> mapping)
+        {
+            var missingHealthStatus = ((HealthStatus[])Enum.GetValues(typeof(HealthStatus))).Except(mapping.Keys).ToList();
+            if (missingHealthStatus.Any())
+            {
+                var missing = string.Join(", ", missingHealthStatus.Select(status => $"{nameof(HealthStatus)}.{status}"));
+                var message =
+                    $"The {nameof(ResultStatusCodes)} dictionary must include an entry for all possible " +
+                    $"{nameof(HealthStatus)} values. Missing: {missing}";
+                throw new InvalidOperationException(message);
+            }
+            return mapping;
+        }
+
         /// <summary>
         /// Gets or sets a delegate used to write the response.
         /// </summary>
diff --git a/src/Middleware/HealthChecks/test/UnitTests/HealthCheckMiddlewareTests.cs b/src/Middleware/HealthChecks/test/UnitTests/HealthCheckMiddlewareTests.cs
index 9bcefafbaa..135a4bed9c 100644
--- a/src/Middleware/HealthChecks/test/UnitTests/HealthCheckMiddlewareTests.cs
+++ b/src/Middleware/HealthChecks/test/UnitTests/HealthCheckMiddlewareTests.cs
@@ -1,6 +1,8 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
+using System;
+using System.Collections.Generic;
 using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Http;
@@ -684,5 +686,68 @@ namespace Microsoft.AspNetCore.Diagnostics.HealthChecks
             Assert.Equal("text/plain", response.Content.Headers.ContentType.ToString());
             Assert.Equal("Healthy", await response.Content.ReadAsStringAsync());
         }
+        
+        [Fact]
+        public void HealthCheckOptions_HasDefaultMappingWithDefaultResultStatusCodes()
+        {
+            var options = new HealthCheckOptions();
+            Assert.NotNull(options.ResultStatusCodes);
+            Assert.Equal(StatusCodes.Status200OK, options.ResultStatusCodes[HealthStatus.Healthy]);
+            Assert.Equal(StatusCodes.Status200OK, options.ResultStatusCodes[HealthStatus.Degraded]);
+            Assert.Equal(StatusCodes.Status503ServiceUnavailable, options.ResultStatusCodes[HealthStatus.Unhealthy]);
+        }
+        
+        [Fact]
+        public void HealthCheckOptions_HasDefaultMappingWhenResettingResultStatusCodes()
+        {
+            var options = new HealthCheckOptions { ResultStatusCodes = null };
+            Assert.NotNull(options.ResultStatusCodes);
+            Assert.Equal(StatusCodes.Status200OK, options.ResultStatusCodes[HealthStatus.Healthy]);
+            Assert.Equal(StatusCodes.Status200OK, options.ResultStatusCodes[HealthStatus.Degraded]);
+            Assert.Equal(StatusCodes.Status503ServiceUnavailable, options.ResultStatusCodes[HealthStatus.Unhealthy]);
+        }
+        
+        
+        [Fact]
+        public void HealthCheckOptions_DoesNotThrowWhenProperlyConfiguringResultStatusCodes()
+        {
+            _ = new HealthCheckOptions
+            {
+                ResultStatusCodes = new Dictionary<HealthStatus, int>
+                {
+                    [HealthStatus.Healthy] = 200,
+                    [HealthStatus.Degraded] = 200,
+                    [HealthStatus.Unhealthy] = 503
+                }
+            };
+        }
+
+        [Fact]
+        public void HealthCheckOptions_ThrowsWhenAHealthStatusIsMissing()
+        {
+            var exception = Assert.Throws<InvalidOperationException>(() => 
+                new HealthCheckOptions { ResultStatusCodes = new Dictionary<HealthStatus, int>() }
+            );
+            Assert.Contains($"{nameof(HealthStatus)}.{nameof(HealthStatus.Healthy)}", exception.Message);
+            Assert.Contains($"{nameof(HealthStatus)}.{nameof(HealthStatus.Degraded)}", exception.Message);
+            Assert.Contains($"{nameof(HealthStatus)}.{nameof(HealthStatus.Unhealthy)}", exception.Message);
+        }
+        
+        [Fact]
+        public void HealthCheckOptions_ThrowsWhenAHealthStatusIsMissing_MessageDoesNotContainDefinedStatus()
+        {
+            var exception = Assert.Throws<InvalidOperationException>(() => 
+                new HealthCheckOptions
+                {
+                    ResultStatusCodes = new Dictionary<HealthStatus, int>
+                    {
+                        [HealthStatus.Healthy] = 200
+                    }
+                }
+            );
+            Assert.DoesNotContain($"{nameof(HealthStatus)}.{nameof(HealthStatus.Healthy)}", exception.Message);
+            Assert.Contains($"{nameof(HealthStatus)}.{nameof(HealthStatus.Degraded)}", exception.Message);
+            Assert.Contains($"{nameof(HealthStatus)}.{nameof(HealthStatus.Unhealthy)}", exception.Message);
+        }
     }
 }