diff --git a/src/Microsoft.AspNet.FeatureModel/FeatureCollection.cs b/src/Microsoft.AspNet.FeatureModel/FeatureCollection.cs index 222dc181f0..3b3f05b2b2 100644 --- a/src/Microsoft.AspNet.FeatureModel/FeatureCollection.cs +++ b/src/Microsoft.AspNet.FeatureModel/FeatureCollection.cs @@ -66,7 +66,12 @@ namespace Microsoft.AspNet.FeatureModel void SetInterface(Type type, object feature) { if (type == null) throw new ArgumentNullException("type"); - if (feature == null) throw new ArgumentNullException("feature"); + + if (feature == null) + { + Remove(type); + return; + } lock (_containerSync) { @@ -167,7 +172,20 @@ namespace Microsoft.AspNet.FeatureModel public bool Remove(Type key) { - throw new NotImplementedException(); + if (key == null) throw new ArgumentNullException("key"); + + lock (_containerSync) + { + Type priorFeatureType; + if (_featureTypeByName.TryGetValue(key.FullName, out priorFeatureType)) + { + _featureTypeByName.Remove(key.FullName); + _featureByFeatureType.Remove(priorFeatureType); + Interlocked.Increment(ref _containerRevision); + return true; + } + return false; + } } public bool TryGetValue(Type key, out object value) diff --git a/test/Microsoft.AspNet.FeatureModel.Tests/InterfaceDictionaryTests.cs b/test/Microsoft.AspNet.FeatureModel.Tests/InterfaceDictionaryTests.cs index 37eac34813..9309814aa7 100644 --- a/test/Microsoft.AspNet.FeatureModel.Tests/InterfaceDictionaryTests.cs +++ b/test/Microsoft.AspNet.FeatureModel.Tests/InterfaceDictionaryTests.cs @@ -48,5 +48,36 @@ namespace Microsoft.AspNet.FeatureModel.Tests Assert.Throws(() => interfaces.Add(typeof(IThing), thing)); } + + [Fact] + public void RemovedInterfaceIsRemoved() + { + var interfaces = new FeatureCollection(); + var thing = new Thing(); + + interfaces.Add(typeof(IThing), thing); + + Assert.Equal(interfaces[typeof(IThing)], thing); + + Assert.True(interfaces.Remove(typeof(IThing))); + + object thing2; + Assert.False(interfaces.TryGetValue(typeof(IThing), out thing2)); + } + + [Fact] + public void SetNullValueRemoves() + { + var interfaces = new FeatureCollection(); + var thing = new Thing(); + + interfaces.Add(typeof(IThing), thing); + Assert.Equal(interfaces[typeof(IThing)], thing); + + interfaces[typeof(IThing)] = null; + + object thing2; + Assert.False(interfaces.TryGetValue(typeof(IThing), out thing2)); + } } }