Add .Unset method to mock

README.md

@@ -190,6 +190,31 @@ func TestSomethingElse(t *testing.T) {
 
 
 }
+
+// TestSomethingElse2 is a third example that shows how you can use
+// the Off method to cleanup handlers and then add new ones.
+func TestSomethingElse2(t *testing.T) {
+
+  // create an instance of our test object
+  testObj := new(MyMockedObject)
+
+  // setup expectations with a placeholder in the argument list
+  testObj.On("DoSomething", mock.Anything).Return(true, nil)
+
+  // call the code we are testing
+  targetFuncThatDoesSomethingWithObj(testObj)
+
+  // assert that the expectations were met
+  testObj.AssertExpectations(t)
+
+  // remove the handler now so we can add another one that takes precedence
+  testObj.Off("DoSomething", mock.Anything)
+
+  // return false now instead of true
+  testObj.On("DoSomething", mock.Anything).Return(false, nil)
+
+  testObj.AssertExpectations(t)
+}
 ```
 
 For more information on how to write mock code, check out the [API documentation for the `mock` package](http://godoc.org/github.com/stretchr/testify/mock).

mock/mock.go

@@ -275,6 +275,42 @@ func (m *Mock) On(methodName string, arguments ...interface{}) *Call {
 	return c
 }
 
+// Off removes a mock handler from being called. You must pass the same exact
+// arguments that were called in the original .On call.
+//
+//     Mock.Off("MyMethod", arg1, arg2)
+func (m *Mock) Off(methodName string, arguments ...interface{}) *Mock {
+	for _, arg := range arguments {
+		if v := reflect.ValueOf(arg); v.Kind() == reflect.Func {
+			panic(fmt.Sprintf("cannot use Func in expectations. Use mock.AnythingOfType(\"%T\")", arg))
+		}
+	}
+
+	m.mutex.Lock()
+
+	foundMatchingCall := false
+
+	for i, call := range m.ExpectedCalls {
+		if call.Method == methodName {
+			_, diffCount := call.Arguments.Diff(arguments)
+			if diffCount == 0 {
+				foundMatchingCall = true
+				// Remove from ExpectedCalls
+				m.ExpectedCalls = append(m.ExpectedCalls[:i], m.ExpectedCalls[i+1:]...)
+			}
+		}
+	}
+
+	m.mutex.Unlock()
+	if !foundMatchingCall {
+		m.fail("\n\nmock: Could not find expected call\n-----------------------------\n\n%s\n\n",
+			callString(methodName, arguments, true),
+		)
+	}
+
+	return m
+}
+
 // /*
 // 	Recording and responding to activity
 // */

mock/mock_test.go

@@ -462,6 +462,61 @@ func Test_Mock_On_WithFuncTypeArg(t *testing.T) {
 	})
 }
 
+func Test_Mock_Off(t *testing.T) {
+	// make a test impl object
+	var mockedService = new(TestExampleImplementation)
+
+	call := mockedService.
+		On("TheExampleMethodFuncType", "argA").
+		Return(nil)
+
+	found, foundCall := mockedService.findExpectedCall("TheExampleMethodFuncType", "argA")
+	require.NotEqual(t, -1, found)
+	require.Equal(t, foundCall, call)
+
+	mockedService.
+		Off("TheExampleMethodFuncType", "argA")
+
+	found, foundCall = mockedService.findExpectedCall("TheExampleMethodFuncType", "argA")
+	require.Equal(t, -1, found)
+
+	var expectedCall *Call
+	require.Equal(t, expectedCall, foundCall)
+
+	fn := func(string) error { return nil }
+	assert.Panics(t, func() {
+		mockedService.TheExampleMethodFuncType(fn)
+	})
+}
+
+func Test_Mock_Chained_Off(t *testing.T) {
+	// make a test impl object
+	var mockedService = new(TestExampleImplementation)
+
+	// determine our current line number so we can assert the expected calls callerInfo properly
+	_, _, line, _ := runtime.Caller(0)
+	mockedService.
+		On("TheExampleMethod1", 1, 1).
+		Return(0).
+		On("TheExampleMethod2", 2, 2).
+		On("TheExampleMethod3", 3, 3, 3).
+		Return(nil)
+
+	mockedService.
+		Off("TheExampleMethod2", 2, 2).
+		Off("TheExampleMethod3", 3, 3, 3)
+
+	expectedCalls := []*Call{
+		{
+			Parent:          &mockedService.Mock,
+			Method:          "TheExampleMethod1",
+			Arguments:       []interface{}{1, 1},
+			ReturnArguments: []interface{}{0},
+			callerInfo:      []string{fmt.Sprintf("mock_test.go:%d", line+2)},
+		},
+	}
+	assert.Equal(t, expectedCalls, mockedService.ExpectedCalls)
+}
 func Test_Mock_Return(t *testing.T) {
 
 	// make a test impl object