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