Add NullRedactor to the RedactorProvider during initialization (#5424)

* Add NullRedactor to the RedactorProvider during initialization. Fix #5265 --------- Co-authored-by: Yifan Zhu <yifzhu@microsoft.com>
dotnet · Sep 19, 2024 · 1d9d44f · 1d9d44f
1 parent 539f3a3
commit 1d9d44f
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 6 deletions.
diff --git a/src/Libraries/Microsoft.Extensions.Compliance.Redaction/RedactorProvider.cs b/src/Libraries/Microsoft.Extensions.Compliance.Redaction/RedactorProvider.cs
@@ -37,6 +37,12 @@ public Redactor GetRedactor(DataClassificationSet classifications)
 
     private static FrozenDictionary<DataClassificationSet, Redactor> GetClassRedactorMap(IEnumerable<Redactor> redactors, Dictionary<DataClassificationSet, Type> map)
     {
+        if (!map.ContainsKey(DataClassification.None))
+        {
+            map.Add(DataClassification.None, typeof(NullRedactor));
+            redactors = [.. redactors, NullRedactor.Instance];
+        }
+
         var dict = new Dictionary<DataClassificationSet, Redactor>(map.Count);
         foreach (var m in map)
         {
@@ -45,6 +51,7 @@ private static FrozenDictionary<DataClassificationSet, Redactor> GetClassRedacto
                 if (r.GetType() == m.Value)
                 {
                     dict[m.Key] = r;
+                    break;
                 }
             }
         }

diff --git a/test/Libraries/Microsoft.Extensions.Compliance.Redaction.Tests/RedactorProviderTests.cs b/test/Libraries/Microsoft.Extensions.Compliance.Redaction.Tests/RedactorProviderTests.cs
@@ -10,6 +10,16 @@ namespace Microsoft.Extensions.Compliance.Redaction.Test;
 
 public class RedactorProviderTests
 {
+    [Fact]
+    public void RedactorProvider_Returns_NullRedactor_For_NoneDataClassification()
+    {
+        var redactorProvider = new RedactorProvider(
+            redactors: [ErasingRedactor.Instance],
+            options: Microsoft.Extensions.Options.Options.Create(new RedactorProviderOptions()));
+
+        Assert.IsType<NullRedactor>(redactorProvider.GetRedactor(DataClassification.None));
+    }
+
     [Fact]
     public void RedactorProvider_Returns_Redactor_For_Every_Data_Classification()
     {
@@ -38,13 +48,15 @@ public void RedactorProvider_Returns_Redactor_For_Data_Classifications()
             redactors: new Redactor[] { ErasingRedactor.Instance, NullRedactor.Instance },
             options: Microsoft.Extensions.Options.Options.Create(opt));
 
-        var r1 = redactorProvider.GetRedactor(_dataClassification1);
-        var r2 = redactorProvider.GetRedactor(_dataClassification2);
-        var r3 = redactorProvider.GetRedactor(_dataClassification3);
+        Redactor r1 = redactorProvider.GetRedactor(_dataClassification1);
+        Redactor r2 = redactorProvider.GetRedactor(_dataClassification2);
+        Redactor r3 = redactorProvider.GetRedactor(_dataClassification3);
+        Redactor r4 = redactorProvider.GetRedactor(DataClassification.None);
 
-        Assert.Equal(typeof(ErasingRedactor), r1.GetType());
-        Assert.Equal(typeof(NullRedactor), r2.GetType());
-        Assert.Equal(typeof(ErasingRedactor), r3.GetType());
+        Assert.IsType<ErasingRedactor>(r1);
+        Assert.IsType<NullRedactor>(r2);
+        Assert.IsType<ErasingRedactor>(r3);
+        Assert.IsType<NullRedactor>(r4);
     }
 
     [Fact]