Feat: Add AddOrUpdate extension methods for dictionaries with tests

Author: AnnaSasDev

src/CodeOfChaos.Extensions/DictionaryExtensions.cs

@@ -18,6 +18,32 @@ TValue value
         dictionary[key] = value;
         return dictionary;
     }
+    
+    public static Dictionary<TKey, TValue> AddOrUpdate<TKey, TValue>(
+        this Dictionary<TKey, TValue> dictionary,
+        TKey key,
+        Func<TKey, TValue> valueFactory
+    ) where TKey : notnull {
+        if (dictionary.TryAdd(key, valueFactory(key))) return dictionary;
+        dictionary[key] = valueFactory(key);
+        return dictionary;
+    }
+    
+    public static Dictionary<TKey, TValue> AddOrUpdate<TKey, TValue>(
+        this Dictionary<TKey, TValue> dictionary,
+        TKey key,
+        Func<TKey, TValue> addValueFactory,
+        Func<TKey, TValue, TValue> updateValueFactory
+    ) where TKey : notnull {
+        if (dictionary.TryGetValue(key, out TValue? value)) {
+            value = updateValueFactory(key, value);
+            dictionary[key] = value;
+            return dictionary;
+        }
+        value = addValueFactory(key);
+        dictionary.Add(key, value);
+        return dictionary;
+    }
 
     public static bool TryAddToOrCreateCollection<TKey, TValue, TCollection>(
         this Dictionary<TKey, TCollection> dictionary,

tests/Tests.CodeOfChaos.Extensions/DictionaryExtensionTests.cs

@@ -109,4 +109,63 @@ public async Task GetOrAdd_ShouldAddValueToDictionary_WhenKeyDoesNotExist_Overlo
         await Assert.That(dictionary.ContainsKey("key1")).IsTrue();
         await Assert.That(dictionary["key1"]).IsEqualTo("value10");
     }
+
+    // ReSharper disable once UnusedParameter.Local
+    // ReSharper disable once ConvertToLocalFunction
+    [Test]
+    public async Task AddOrUpdate_ShouldAddValueToDictionary_WhenKeyDoesNotExist_OverloadFactory() {
+        // Arrange
+        var dictionary = new Dictionary<string, string>();
+        const string value = "value";
+        Func<string, string> valueFactory = key => value;
+
+        // Act
+        dictionary.AddOrUpdate("key1", valueFactory);
+
+        // Assert
+        await Assert.That(dictionary.ContainsKey("key1")).IsTrue();
+        await Assert.That(dictionary["key1"]).IsEqualTo(value);
+    }
+
+    // ReSharper disable once UnusedParameter.Local
+    // ReSharper disable once ConvertToLocalFunction
+    [Test]
+    public async Task AddOrUpdate_ShouldAddValueToDictionary_WhenKeyDoesExist_OverloadFactory() {
+        // Arrange
+        var dictionary = new Dictionary<string, string> { ["key1"] = "oldValue" };
+
+        const string value = "value";
+        Func<string, string> valueFactory = key => $"{value}";
+
+        // Act
+        dictionary.AddOrUpdate("key1", valueFactory);
+
+        // Assert
+        await Assert.That(dictionary).ContainsKey("key1");
+        await Assert.That(dictionary["key1"])
+            .IsNotEqualTo("oldValue")
+            .IsEqualTo("value");
+    }
+    
+    // ReSharper disable once UnusedParameter.Local
+    // ReSharper disable once ConvertToLocalFunction
+    [Test]
+    public async Task AddOrUpdate_ShouldAddValueToDictionary_WhenKeyDoesNotExist_OverloadFactoryAddAndUpdate() {
+        // Arrange
+        var dictionary = new Dictionary<string, string> { ["key1"] = "oldValue" };
+        const string value = "value";
+        Func<string, string> onAddFactory = key => value;
+        Func<string, string, string> onUpdateFactory = (_, oldValue) => $"{oldValue};{value}";
+
+        // Act
+        dictionary.AddOrUpdate("key1", onAddFactory, onUpdateFactory);
+        dictionary.AddOrUpdate("key2", onAddFactory, onUpdateFactory);
+
+        // Assert
+        await Assert.That(dictionary).ContainsKey("key1");
+        await Assert.That(dictionary["key1"]).IsEqualTo("oldValue;value");
+        
+        await Assert.That(dictionary).ContainsKey("key2");
+        await Assert.That(dictionary["key2"]).IsEqualTo(value);
+    }
 }