Restore provider correctly in FluentMockContext

devlooped · Dec 3, 2017 · 7ffb12e · 7ffb12e
1 parent 2818672
commit 7ffb12e
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/Source/FluentMockContext.cs b/Source/FluentMockContext.cs
@@ -100,14 +100,15 @@ public void Dispose()
 
 		internal class MockInvocation : IDisposable
 		{
-			private DefaultValue defaultValue;
+			private DefaultValueProvider defaultValueProvider;
 
 			public MockInvocation(Mock mock, Invocation invocation, Match matcher)
 			{
 				this.Mock = mock;
 				this.Invocation = invocation;
 				this.Match = matcher;
-				defaultValue = mock.DefaultValue;
+				this.defaultValueProvider = mock.DefaultValueProvider;
+
 				// Temporarily set mock default value to Mock so that recursion works.
 				mock.DefaultValue = DefaultValue.Mock;
 			}
@@ -120,7 +121,7 @@ public MockInvocation(Mock mock, Invocation invocation, Match matcher)
 
 			public void Dispose()
 			{
-				Mock.DefaultValue = defaultValue;
+				this.Mock.DefaultValueProvider = this.defaultValueProvider;
 			}
 		}
 	}