Merge pull request #158 from stretchr/mock-run-handler

Allow mocking methods updating referenced structs

mock/mock.go

@@ -42,6 +42,11 @@ type Call struct {
 	// Holds a channel that will be used to block the Return until it either
 	// recieves a message or is closed. nil means it returns immediately.
 	WaitFor <-chan time.Time
+
+	// Holds a handler used to manipulate arguments content that are passed by
+	// reference. It's useful when mocking methods such as unmarshalers or
+	// decoders.
+	Run func(Arguments)
 }
 
 // Mock is the workhorse used to track activity on another object.
@@ -100,7 +105,7 @@ func (m *Mock) On(methodName string, arguments ...interface{}) *Mock {
 //
 //     Mock.On("MyMethod", arg1, arg2).Return(returnArg1, returnArg2)
 func (m *Mock) Return(returnArguments ...interface{}) *Mock {
-	m.ExpectedCalls = append(m.ExpectedCalls, Call{m.onMethodName, m.onMethodArguments, returnArguments, 0, nil})
+	m.ExpectedCalls = append(m.ExpectedCalls, Call{m.onMethodName, m.onMethodArguments, returnArguments, 0, nil, nil})
 	return m
 }
 
@@ -142,6 +147,19 @@ func (m *Mock) After(d time.Duration) *Mock {
 	return m.WaitUntil(time.After(d))
 }
 
+// Run sets a handler to be called before returning. It can be used when
+// mocking a method such as unmarshalers that takes a pointer to a struct and
+// sets properties in such struct
+//
+//    Mock.On("Unmarshal", AnythingOfType("*map[string]interface{}").Return().Run(function(args Arguments) {
+//    	arg := args.Get(0).(*map[string]interface{})
+//    	arg["foo"] = "bar"
+//    })
+func (m *Mock) Run(fn func(Arguments)) *Mock {
+	m.ExpectedCalls[len(m.ExpectedCalls)-1].Run = fn
+	return m
+}
+
 /*
 	Recording and responding to activity
 */
@@ -242,13 +260,17 @@ func (m *Mock) Called(arguments ...interface{}) Arguments {
 	}
 
 	// add the call
-	m.Calls = append(m.Calls, Call{functionName, arguments, make([]interface{}, 0), 0, nil})
+	m.Calls = append(m.Calls, Call{functionName, arguments, make([]interface{}, 0), 0, nil, nil})
 
 	// block if specified
 	if call.WaitFor != nil {
 		<-call.WaitFor
 	}
 
+	if call.Run != nil {
+		call.Run(arguments)
+	}
+
 	return call.ReturnArguments
 
 }

mock/mock_test.go

@@ -30,7 +30,9 @@ func (i *TestExampleImplementation) TheExampleMethod2(yesorno bool) {
 	i.Called(yesorno)
 }
 
-type ExampleType struct{}
+type ExampleType struct {
+	ran bool
+}
 
 func (i *TestExampleImplementation) TheExampleMethod3(et *ExampleType) error {
 	args := i.Called(et)
@@ -153,6 +155,36 @@ func Test_Mock_Return_After(t *testing.T) {
 
 }
 
+func Test_Mock_Return_Run(t *testing.T) {
+
+	// make a test impl object
+	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
+
+	assert.Equal(t, mockedService.Mock.On("TheExampleMethod3", AnythingOfType("*mock.ExampleType")).Return(nil).Run(func(args Arguments) {
+		arg := args.Get(0).(*ExampleType)
+		arg.ran = true
+	}), &mockedService.Mock)
+
+	// ensure the call was created
+	if assert.Equal(t, 1, len(mockedService.Mock.ExpectedCalls)) {
+		call := mockedService.Mock.ExpectedCalls[0]
+
+		assert.Equal(t, "TheExampleMethod3", call.Method)
+		assert.Equal(t, AnythingOfType("*mock.ExampleType"), call.Arguments[0])
+		assert.Equal(t, nil, call.ReturnArguments[0])
+		assert.Equal(t, 0, call.Repeatability)
+		assert.NotEqual(t, nil, call.WaitFor)
+		assert.NotNil(t, call.Run)
+
+	}
+
+	et := ExampleType{}
+	assert.Equal(t, false, et.ran)
+	mockedService.TheExampleMethod3(&et)
+	assert.Equal(t, true, et.ran)
+
+}
+
 func Test_Mock_Return_Once(t *testing.T) {
 
 	// make a test impl object