Allows mock call expectations to be defined in arbitrary order

Originally the expctations required that the return values be specified
immediately after the method name and arguments, otherwise the call
setup will either panic (best case) or silently modify the *previous*
call specification (worst case).

This change moves the Return(), Run(), Once(), etc methods onto the Call
struct, and changes the chaining behaviour so that they modify the Call
data directly rather than referencing the last item in the ExpectedCalls
array.

mock/mock.go

@@ -25,6 +25,7 @@ type TestingT interface {
 // Call represents a method call and is used for setting expectations,
 // as well as recording activity.
 type Call struct {
+	Parent *Mock
 
 	// The name of the method that was or will be called.
 	Method string
@@ -47,24 +48,113 @@ type Call struct {
 	// 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)
+	RunFn func(Arguments)
+}
+
+func newCall(parent *Mock, methodName string, methodArguments ...interface{}) *Call {
+	return &Call{
+		Parent:          parent,
+		Method:          methodName,
+		Arguments:       methodArguments,
+		ReturnArguments: make([]interface{}, 0),
+		Repeatability:   0,
+		WaitFor:         nil,
+		RunFn:           nil,
+	}
+}
+
+func (self *Call) lock() {
+	self.Parent.mutex.Lock()
+}
+
+func (self *Call) unlock() {
+	self.Parent.mutex.Unlock()
+}
+
+func (self *Call) Return(returnArguments ...interface{}) *Call {
+	self.lock()
+	defer self.unlock()
+
+	self.ReturnArguments = returnArguments
+
+	return self
+}
+
+// Once indicates that that the mock should only return the value once.
+//
+//    Mock.On("MyMethod", arg1, arg2).Return(returnArg1, returnArg2).Once()
+func (self *Call) Once() *Call {
+	return self.Times(1)
+}
+
+// Twice indicates that that the mock should only return the value twice.
+//
+//    Mock.On("MyMethod", arg1, arg2).Return(returnArg1, returnArg2).Twice()
+func (self *Call) Twice() *Call {
+	return self.Times(2)
+}
+
+// Times indicates that that the mock should only return the indicated number
+// of times.
+//
+//    Mock.On("MyMethod", arg1, arg2).Return(returnArg1, returnArg2).Times(5)
+func (self *Call) Times(i int) *Call {
+	self.lock()
+	defer self.unlock()
+	self.Repeatability = i
+	return self
+}
+
+// WaitUntil sets the channel that will block the mock's return until its closed
+// or a message is received.
+//
+//    Mock.On("MyMethod", arg1, arg2).WaitUntil(time.After(time.Second))
+func (self *Call) WaitUntil(w <-chan time.Time) *Call {
+	self.lock()
+	defer self.unlock()
+	self.WaitFor = w
+	return self
+}
+
+// After sets how long to block until the call returns
+//
+//    Mock.On("MyMethod", arg1, arg2).After(time.Second)
+func (self *Call) After(d time.Duration) *Call {
+	return self.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 (self *Call) Run(fn func(Arguments)) *Call {
+	self.lock()
+	defer self.unlock()
+	self.RunFn = fn
+	return self
+}
+
+// On chains a new expectation description onto the mocked interface. This
+// allows syntax like.
+//
+//    Mock.
+//       On("MyMethod", 1).Return(nil).
+//       On("MyOtherMethod", 'a', 'b', 'c').Return(errors.New("Some Error"))
+func (self *Call) On(methodName string, arguments ...interface{}) *Call {
+	return self.Parent.On(methodName, arguments...)
 }
 
 // Mock is the workhorse used to track activity on another object.
-// For an example of its usage, refer to the "Example Usage" section at the top of this document.
+// For an example of its usage, refer to the "Example Usage" section at the top
+// of this document.
 type Mock struct {
-
-	// The method name that is currently
-	// being referred to by the On method.
-	onMethodName string
-
-	// An array of the arguments that are
-	// currently being referred to by the On method.
-	onMethodArguments Arguments
-
 	// Represents the calls that are expected of
 	// an object.
-	ExpectedCalls []Call
+	ExpectedCalls []*Call
 
 	// Holds the calls that were made to this mocked object.
 	Calls []Call
@@ -95,95 +185,23 @@ func (m *Mock) TestData() objx.Map {
 // being called.
 //
 //     Mock.On("MyMethod", arg1, arg2)
-func (m *Mock) On(methodName string, arguments ...interface{}) *Mock {
-	m.onMethodName = methodName
-	m.onMethodArguments = arguments
-
+func (self *Mock) On(methodName string, arguments ...interface{}) *Call {
 	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))
 		}
 	}
 
-	return m
+	self.mutex.Lock()
+	defer self.mutex.Unlock()
+	c := newCall(self, methodName, arguments...)
+	self.ExpectedCalls = append(self.ExpectedCalls, c)
+	return c
 }
 
-// Return finishes a description of an expectation of the method (and arguments)
-// specified in the most recent On method call.
-//
-//     Mock.On("MyMethod", arg1, arg2).Return(returnArg1, returnArg2)
-func (m *Mock) Return(returnArguments ...interface{}) *Mock {
-	m.mutex.Lock()
-	defer m.mutex.Unlock()
-
-	m.ExpectedCalls = append(m.ExpectedCalls, Call{m.onMethodName, m.onMethodArguments, returnArguments, 0, nil, nil})
-	return m
-}
-
-// Once indicates that that the mock should only return the value once.
-//
-//    Mock.On("MyMethod", arg1, arg2).Return(returnArg1, returnArg2).Once()
-func (m *Mock) Once() {
-	m.mutex.Lock()
-	m.ExpectedCalls[len(m.ExpectedCalls)-1].Repeatability = 1
-	m.mutex.Unlock()
-}
-
-// Twice indicates that that the mock should only return the value twice.
-//
-//    Mock.On("MyMethod", arg1, arg2).Return(returnArg1, returnArg2).Twice()
-func (m *Mock) Twice() {
-	m.mutex.Lock()
-	m.ExpectedCalls[len(m.ExpectedCalls)-1].Repeatability = 2
-	m.mutex.Unlock()
-}
-
-// Times indicates that that the mock should only return the indicated number
-// of times.
-//
-//    Mock.On("MyMethod", arg1, arg2).Return(returnArg1, returnArg2).Times(5)
-func (m *Mock) Times(i int) {
-	m.mutex.Lock()
-	m.ExpectedCalls[len(m.ExpectedCalls)-1].Repeatability = i
-	m.mutex.Unlock()
-}
-
-// WaitUntil sets the channel that will block the mock's return until its closed
-// or a message is received.
-//
-//    Mock.On("MyMethod", arg1, arg2).WaitUntil(time.After(time.Second))
-func (m *Mock) WaitUntil(w <-chan time.Time) *Mock {
-	m.mutex.Lock()
-	m.ExpectedCalls[len(m.ExpectedCalls)-1].WaitFor = w
-	m.mutex.Unlock()
-	return m
-}
-
-// After sets how long to block until the call returns
-//
-//    Mock.On("MyMethod", arg1, arg2).After(time.Second)
-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.mutex.Lock()
-	m.ExpectedCalls[len(m.ExpectedCalls)-1].Run = fn
-	m.mutex.Unlock()
-	return m
-}
-
-/*
-	Recording and responding to activity
-*/
+// /*
+// 	Recording and responding to activity
+// */
 
 func (m *Mock) findExpectedCall(method string, arguments ...interface{}) (int, *Call) {
 	for i, call := range m.expectedCalls() {
@@ -191,7 +209,7 @@ func (m *Mock) findExpectedCall(method string, arguments ...interface{}) (int, *
 
 			_, diffCount := call.Arguments.Diff(arguments)
 			if diffCount == 0 {
-				return i, &call
+				return i, call
 			}
 
 		}
@@ -209,7 +227,7 @@ func (m *Mock) findClosestCall(method string, arguments ...interface{}) (bool, *
 			_, tempDiffCount := call.Arguments.Diff(arguments)
 			if tempDiffCount < diffCount || diffCount == 0 {
 				diffCount = tempDiffCount
-				closestCall = &call
+				closestCall = call
 			}
 
 		}
@@ -272,17 +290,16 @@ func (m *Mock) Called(arguments ...interface{}) Arguments {
 		switch {
 		case call.Repeatability == 1:
 			call.Repeatability = -1
-			m.ExpectedCalls[found] = *call
+
 		case call.Repeatability > 1:
 			call.Repeatability -= 1
-			m.ExpectedCalls[found] = *call
 		}
 		m.mutex.Unlock()
 	}
 
 	// add the call
 	m.mutex.Lock()
-	m.Calls = append(m.Calls, Call{functionName, arguments, make([]interface{}, 0), 0, nil, nil})
+	m.Calls = append(m.Calls, *newCall(m, functionName, arguments...))
 	m.mutex.Unlock()
 
 	// block if specified
@@ -290,12 +307,11 @@ func (m *Mock) Called(arguments ...interface{}) Arguments {
 		<-call.WaitFor
 	}
 
-	if call.Run != nil {
-		call.Run(arguments)
+	if call.RunFn != nil {
+		call.RunFn(arguments)
 	}
 
 	return call.ReturnArguments
-
 }
 
 /*
@@ -390,10 +406,10 @@ func (m *Mock) methodWasCalled(methodName string, expected []interface{}) bool {
 	return false
 }
 
-func (m *Mock) expectedCalls() []Call {
+func (m *Mock) expectedCalls() []*Call {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
-	return append([]Call{}, m.ExpectedCalls...)
+	return append([]*Call{}, m.ExpectedCalls...)
 }
 
 func (m *Mock) calls() []Call {

mock/mock_test.go

@@ -3,6 +3,7 @@ package mock
 import (
 	"errors"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"testing"
 	"time"
 )
@@ -73,9 +74,7 @@ func Test_Mock_TestData(t *testing.T) {
 
 		mockedService.TestData().Set("something", 123)
 		assert.Equal(t, 123, mockedService.TestData().Get("something").Data())
-
 	}
-
 }
 
 func Test_Mock_On(t *testing.T) {
@@ -83,9 +82,36 @@ func Test_Mock_On(t *testing.T) {
 	// make a test impl object
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
-	assert.Equal(t, mockedService.On("TheExampleMethod"), &mockedService.Mock)
-	assert.Equal(t, "TheExampleMethod", mockedService.onMethodName)
+	c := mockedService.On("TheExampleMethod")
+	assert.Equal(t, []*Call{c}, mockedService.ExpectedCalls)
+	assert.Equal(t, "TheExampleMethod", c.Method)
+}
 
+func Test_Mock_Chained_On(t *testing.T) {
+	// make a test impl object
+	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
+
+	mockedService.
+		On("TheExampleMethod", 1, 2, 3).
+		Return(0).
+		On("TheExampleMethod3", AnythingOfType("*mock.ExampleType")).
+		Return(nil)
+
+	expectedCalls := []*Call{
+		&Call{
+			Parent:          &mockedService.Mock,
+			Method:          "TheExampleMethod",
+			Arguments:       []interface{}{1, 2, 3},
+			ReturnArguments: []interface{}{0},
+		},
+		&Call{
+			Parent:          &mockedService.Mock,
+			Method:          "TheExampleMethod3",
+			Arguments:       []interface{}{AnythingOfType("*mock.ExampleType")},
+			ReturnArguments: []interface{}{nil},
+		},
+	}
+	assert.Equal(t, expectedCalls, mockedService.ExpectedCalls)
 }
 
 func Test_Mock_On_WithArgs(t *testing.T) {
@@ -93,12 +119,11 @@ func Test_Mock_On_WithArgs(t *testing.T) {
 	// make a test impl object
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
-	assert.Equal(t, mockedService.On("TheExampleMethod", 1, 2, 3), &mockedService.Mock)
-	assert.Equal(t, "TheExampleMethod", mockedService.onMethodName)
-	assert.Equal(t, 1, mockedService.onMethodArguments[0])
-	assert.Equal(t, 2, mockedService.onMethodArguments[1])
-	assert.Equal(t, 3, mockedService.onMethodArguments[2])
+	c := mockedService.On("TheExampleMethod", 1, 2, 3, 4)
 
+	assert.Equal(t, []*Call{c}, mockedService.ExpectedCalls)
+	assert.Equal(t, "TheExampleMethod", c.Method)
+	assert.Equal(t, Arguments{1, 2, 3, 4}, c.Arguments)
 }
 
 func Test_Mock_On_WithFuncArg(t *testing.T) {
@@ -106,13 +131,20 @@ func Test_Mock_On_WithFuncArg(t *testing.T) {
 	// make a test impl object
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
-	assert.Equal(t, mockedService.On("TheExampleMethodFunc", AnythingOfType("func(string) error")).Return(nil), &mockedService.Mock)
-	assert.Equal(t, "TheExampleMethodFunc", mockedService.onMethodName)
-	assert.Equal(t, AnythingOfType("func(string) error"), mockedService.onMethodArguments[0])
+	c := mockedService.
+		On("TheExampleMethodFunc", AnythingOfType("func(string) error")).
+		Return(nil)
+
+	assert.Equal(t, []*Call{c}, mockedService.ExpectedCalls)
+	assert.Equal(t, "TheExampleMethodFunc", c.Method)
+	assert.Equal(t, 1, len(c.Arguments))
+	assert.Equal(t, AnythingOfType("func(string) error"), c.Arguments[0])
 
 	fn := func(string) error { return nil }
-	mockedService.TheExampleMethodFunc(fn)
 
+	assert.NotPanics(t, func() {
+		mockedService.TheExampleMethodFunc(fn)
+	})
 }
 
 func Test_Mock_On_WithVariadicFunc(t *testing.T) {
@@ -120,9 +152,13 @@ func Test_Mock_On_WithVariadicFunc(t *testing.T) {
 	// make a test impl object
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
-	assert.Equal(t, mockedService.On("TheExampleMethodVariadic", []int{1, 2, 3}).Return(nil), &mockedService.Mock)
-	assert.Equal(t, "TheExampleMethodVariadic", mockedService.onMethodName)
-	assert.Equal(t, []int{1, 2, 3}, mockedService.onMethodArguments[0])
+	c := mockedService.
+		On("TheExampleMethodVariadic", []int{1, 2, 3}).
+		Return(nil)
+
+	assert.Equal(t, []*Call{c}, mockedService.ExpectedCalls)
+	assert.Equal(t, 1, len(c.Arguments))
+	assert.Equal(t, []int{1, 2, 3}, c.Arguments[0])
 
 	assert.NotPanics(t, func() {
 		mockedService.TheExampleMethodVariadic(1, 2, 3)
@@ -138,9 +174,12 @@ func Test_Mock_On_WithVariadicFuncWithInterface(t *testing.T) {
 	// make a test impl object
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
-	assert.Equal(t, mockedService.On("TheExampleMethodVariadicInterface", []interface{}{1, 2, 3}).Return(nil), &mockedService.Mock)
-	assert.Equal(t, "TheExampleMethodVariadicInterface", mockedService.onMethodName)
-	assert.Equal(t, []interface{}{1, 2, 3}, mockedService.onMethodArguments[0])
+	c := mockedService.On("TheExampleMethodVariadicInterface", []interface{}{1, 2, 3}).
+		Return(nil)
+
+	assert.Equal(t, []*Call{c}, mockedService.ExpectedCalls)
+	assert.Equal(t, 1, len(c.Arguments))
+	assert.Equal(t, []interface{}{1, 2, 3}, c.Arguments[0])
 
 	assert.NotPanics(t, func() {
 		mockedService.TheExampleMethodVariadicInterface(1, 2, 3)
@@ -157,9 +196,13 @@ func Test_Mock_On_WithVariadicFuncWithEmptyInterfaceArray(t *testing.T) {
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
 	var expected []interface{}
-	assert.Equal(t, mockedService.On("TheExampleMethodVariadicInterface", expected).Return(nil), &mockedService.Mock)
-	assert.Equal(t, "TheExampleMethodVariadicInterface", mockedService.onMethodName)
-	assert.Equal(t, expected, mockedService.onMethodArguments[0])
+	c := mockedService.
+		On("TheExampleMethodVariadicInterface", expected).
+		Return(nil)
+
+	assert.Equal(t, []*Call{c}, mockedService.ExpectedCalls)
+	assert.Equal(t, 1, len(c.Arguments))
+	assert.Equal(t, expected, c.Arguments[0])
 
 	assert.NotPanics(t, func() {
 		mockedService.TheExampleMethodVariadicInterface()
@@ -184,13 +227,18 @@ func Test_Mock_On_WithFuncTypeArg(t *testing.T) {
 	// make a test impl object
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
-	assert.Equal(t, mockedService.On("TheExampleMethodFuncType", AnythingOfType("mock.ExampleFuncType")).Return(nil), &mockedService.Mock)
-	assert.Equal(t, "TheExampleMethodFuncType", mockedService.onMethodName)
-	assert.Equal(t, AnythingOfType("mock.ExampleFuncType"), mockedService.onMethodArguments[0])
+	c := mockedService.
+		On("TheExampleMethodFuncType", AnythingOfType("mock.ExampleFuncType")).
+		Return(nil)
+
+	assert.Equal(t, []*Call{c}, mockedService.ExpectedCalls)
+	assert.Equal(t, 1, len(c.Arguments))
+	assert.Equal(t, AnythingOfType("mock.ExampleFuncType"), c.Arguments[0])
 
 	fn := func(string) error { return nil }
-	mockedService.TheExampleMethodFuncType(fn)
-
+	assert.NotPanics(t, func() {
+		mockedService.TheExampleMethodFuncType(fn)
+	})
 }
 
 func Test_Mock_Return(t *testing.T) {
@@ -198,24 +246,23 @@ func Test_Mock_Return(t *testing.T) {
 	// make a test impl object
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
-	assert.Equal(t, mockedService.On("TheExampleMethod", "A", "B", true).Return(1, "two", true), &mockedService.Mock)
+	c := mockedService.
+		On("TheExampleMethod", "A", "B", true).
+		Return(1, "two", true)
 
-	// ensure the call was created
-	if assert.Equal(t, 1, len(mockedService.ExpectedCalls)) {
-		call := mockedService.ExpectedCalls[0]
+	require.Equal(t, []*Call{c}, mockedService.ExpectedCalls)
 
-		assert.Equal(t, "TheExampleMethod", call.Method)
-		assert.Equal(t, "A", call.Arguments[0])
-		assert.Equal(t, "B", call.Arguments[1])
-		assert.Equal(t, true, call.Arguments[2])
-		assert.Equal(t, 1, call.ReturnArguments[0])
-		assert.Equal(t, "two", call.ReturnArguments[1])
-		assert.Equal(t, true, call.ReturnArguments[2])
-		assert.Equal(t, 0, call.Repeatability)
-		assert.Nil(t, call.WaitFor)
-
-	}
+	call := mockedService.ExpectedCalls[0]
 
+	assert.Equal(t, "TheExampleMethod", call.Method)
+	assert.Equal(t, "A", call.Arguments[0])
+	assert.Equal(t, "B", call.Arguments[1])
+	assert.Equal(t, true, call.Arguments[2])
+	assert.Equal(t, 1, call.ReturnArguments[0])
+	assert.Equal(t, "two", call.ReturnArguments[1])
+	assert.Equal(t, true, call.ReturnArguments[2])
+	assert.Equal(t, 0, call.Repeatability)
+	assert.Nil(t, call.WaitFor)
 }
 
 func Test_Mock_Return_WaitUntil(t *testing.T) {
@@ -224,24 +271,25 @@ func Test_Mock_Return_WaitUntil(t *testing.T) {
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 	ch := time.After(time.Second)
 
-	assert.Equal(t, mockedService.Mock.On("TheExampleMethod", "A", "B", true).Return(1, "two", true).WaitUntil(ch), &mockedService.Mock)
+	c := mockedService.Mock.
+		On("TheExampleMethod", "A", "B", true).
+		WaitUntil(ch).
+		Return(1, "two", true)
 
-	// ensure the call was created
-	if assert.Equal(t, 1, len(mockedService.Mock.ExpectedCalls)) {
-		call := mockedService.Mock.ExpectedCalls[0]
+	// assert that the call was created
+	require.Equal(t, []*Call{c}, mockedService.ExpectedCalls)
 
-		assert.Equal(t, "TheExampleMethod", call.Method)
-		assert.Equal(t, "A", call.Arguments[0])
-		assert.Equal(t, "B", call.Arguments[1])
-		assert.Equal(t, true, call.Arguments[2])
-		assert.Equal(t, 1, call.ReturnArguments[0])
-		assert.Equal(t, "two", call.ReturnArguments[1])
-		assert.Equal(t, true, call.ReturnArguments[2])
-		assert.Equal(t, 0, call.Repeatability)
-		assert.Equal(t, ch, call.WaitFor)
-
-	}
+	call := mockedService.ExpectedCalls[0]
 
+	assert.Equal(t, "TheExampleMethod", call.Method)
+	assert.Equal(t, "A", call.Arguments[0])
+	assert.Equal(t, "B", call.Arguments[1])
+	assert.Equal(t, true, call.Arguments[2])
+	assert.Equal(t, 1, call.ReturnArguments[0])
+	assert.Equal(t, "two", call.ReturnArguments[1])
+	assert.Equal(t, true, call.ReturnArguments[2])
+	assert.Equal(t, 0, call.Repeatability)
+	assert.Equal(t, ch, call.WaitFor)
 }
 
 func Test_Mock_Return_After(t *testing.T) {
@@ -249,23 +297,24 @@ func Test_Mock_Return_After(t *testing.T) {
 	// make a test impl object
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
-	assert.Equal(t, mockedService.Mock.On("TheExampleMethod", "A", "B", true).Return(1, "two", true).After(time.Second), &mockedService.Mock)
+	c := mockedService.Mock.
+		On("TheExampleMethod", "A", "B", true).
+		Return(1, "two", true).
+		After(time.Second)
 
-	// ensure the call was created
-	if assert.Equal(t, 1, len(mockedService.Mock.ExpectedCalls)) {
-		call := mockedService.Mock.ExpectedCalls[0]
+	require.Equal(t, []*Call{c}, mockedService.ExpectedCalls)
 
-		assert.Equal(t, "TheExampleMethod", call.Method)
-		assert.Equal(t, "A", call.Arguments[0])
-		assert.Equal(t, "B", call.Arguments[1])
-		assert.Equal(t, true, call.Arguments[2])
-		assert.Equal(t, 1, call.ReturnArguments[0])
-		assert.Equal(t, "two", call.ReturnArguments[1])
-		assert.Equal(t, true, call.ReturnArguments[2])
-		assert.Equal(t, 0, call.Repeatability)
-		assert.NotEqual(t, nil, call.WaitFor)
+	call := mockedService.Mock.ExpectedCalls[0]
 
-	}
+	assert.Equal(t, "TheExampleMethod", call.Method)
+	assert.Equal(t, "A", call.Arguments[0])
+	assert.Equal(t, "B", call.Arguments[1])
+	assert.Equal(t, true, call.Arguments[2])
+	assert.Equal(t, 1, call.ReturnArguments[0])
+	assert.Equal(t, "two", call.ReturnArguments[1])
+	assert.Equal(t, true, call.ReturnArguments[2])
+	assert.Equal(t, 0, call.Repeatability)
+	assert.NotEqual(t, nil, call.WaitFor)
 
 }
 
@@ -274,29 +323,56 @@ 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) {
+	fn := 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)
-
 	}
 
+	c := mockedService.Mock.
+		On("TheExampleMethod3", AnythingOfType("*mock.ExampleType")).
+		Return(nil).
+		Run(fn)
+
+	require.Equal(t, []*Call{c}, mockedService.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_Run_Out_Of_Order(t *testing.T) {
+	// make a test impl object
+	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
+	f := func(args Arguments) {
+		arg := args.Get(0).(*ExampleType)
+		arg.ran = true
+	}
+
+	c := mockedService.Mock.
+		On("TheExampleMethod3", AnythingOfType("*mock.ExampleType")).
+		Run(f).
+		Return(nil)
+
+	require.Equal(t, []*Call{c}, mockedService.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)
 }
 
 func Test_Mock_Return_Once(t *testing.T) {
@@ -304,24 +380,23 @@ func Test_Mock_Return_Once(t *testing.T) {
 	// make a test impl object
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
-	mockedService.On("TheExampleMethod", "A", "B", true).Return(1, "two", true).Once()
+	c := mockedService.On("TheExampleMethod", "A", "B", true).
+		Return(1, "two", true).
+		Once()
 
-	// ensure the call was created
-	if assert.Equal(t, 1, len(mockedService.ExpectedCalls)) {
-		call := mockedService.ExpectedCalls[0]
+	require.Equal(t, []*Call{c}, mockedService.ExpectedCalls)
 
-		assert.Equal(t, "TheExampleMethod", call.Method)
-		assert.Equal(t, "A", call.Arguments[0])
-		assert.Equal(t, "B", call.Arguments[1])
-		assert.Equal(t, true, call.Arguments[2])
-		assert.Equal(t, 1, call.ReturnArguments[0])
-		assert.Equal(t, "two", call.ReturnArguments[1])
-		assert.Equal(t, true, call.ReturnArguments[2])
-		assert.Equal(t, 1, call.Repeatability)
-		assert.Nil(t, call.WaitFor)
-
-	}
+	call := mockedService.ExpectedCalls[0]
 
+	assert.Equal(t, "TheExampleMethod", call.Method)
+	assert.Equal(t, "A", call.Arguments[0])
+	assert.Equal(t, "B", call.Arguments[1])
+	assert.Equal(t, true, call.Arguments[2])
+	assert.Equal(t, 1, call.ReturnArguments[0])
+	assert.Equal(t, "two", call.ReturnArguments[1])
+	assert.Equal(t, true, call.ReturnArguments[2])
+	assert.Equal(t, 1, call.Repeatability)
+	assert.Nil(t, call.WaitFor)
 }
 
 func Test_Mock_Return_Twice(t *testing.T) {
@@ -329,24 +404,24 @@ func Test_Mock_Return_Twice(t *testing.T) {
 	// make a test impl object
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
-	mockedService.On("TheExampleMethod", "A", "B", true).Return(1, "two", true).Twice()
+	c := mockedService.
+		On("TheExampleMethod", "A", "B", true).
+		Return(1, "two", true).
+		Twice()
 
-	// ensure the call was created
-	if assert.Equal(t, 1, len(mockedService.ExpectedCalls)) {
-		call := mockedService.ExpectedCalls[0]
+	require.Equal(t, []*Call{c}, mockedService.ExpectedCalls)
 
-		assert.Equal(t, "TheExampleMethod", call.Method)
-		assert.Equal(t, "A", call.Arguments[0])
-		assert.Equal(t, "B", call.Arguments[1])
-		assert.Equal(t, true, call.Arguments[2])
-		assert.Equal(t, 1, call.ReturnArguments[0])
-		assert.Equal(t, "two", call.ReturnArguments[1])
-		assert.Equal(t, true, call.ReturnArguments[2])
-		assert.Equal(t, 2, call.Repeatability)
-		assert.Nil(t, call.WaitFor)
-
-	}
+	call := mockedService.ExpectedCalls[0]
 
+	assert.Equal(t, "TheExampleMethod", call.Method)
+	assert.Equal(t, "A", call.Arguments[0])
+	assert.Equal(t, "B", call.Arguments[1])
+	assert.Equal(t, true, call.Arguments[2])
+	assert.Equal(t, 1, call.ReturnArguments[0])
+	assert.Equal(t, "two", call.ReturnArguments[1])
+	assert.Equal(t, true, call.ReturnArguments[2])
+	assert.Equal(t, 2, call.Repeatability)
+	assert.Nil(t, call.WaitFor)
 }
 
 func Test_Mock_Return_Times(t *testing.T) {
@@ -354,24 +429,24 @@ func Test_Mock_Return_Times(t *testing.T) {
 	// make a test impl object
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
-	mockedService.On("TheExampleMethod", "A", "B", true).Return(1, "two", true).Times(5)
+	c := mockedService.
+		On("TheExampleMethod", "A", "B", true).
+		Return(1, "two", true).
+		Times(5)
 
-	// ensure the call was created
-	if assert.Equal(t, 1, len(mockedService.ExpectedCalls)) {
-		call := mockedService.ExpectedCalls[0]
+	require.Equal(t, []*Call{c}, mockedService.ExpectedCalls)
 
-		assert.Equal(t, "TheExampleMethod", call.Method)
-		assert.Equal(t, "A", call.Arguments[0])
-		assert.Equal(t, "B", call.Arguments[1])
-		assert.Equal(t, true, call.Arguments[2])
-		assert.Equal(t, 1, call.ReturnArguments[0])
-		assert.Equal(t, "two", call.ReturnArguments[1])
-		assert.Equal(t, true, call.ReturnArguments[2])
-		assert.Equal(t, 5, call.Repeatability)
-		assert.Nil(t, call.WaitFor)
-
-	}
+	call := mockedService.ExpectedCalls[0]
 
+	assert.Equal(t, "TheExampleMethod", call.Method)
+	assert.Equal(t, "A", call.Arguments[0])
+	assert.Equal(t, "B", call.Arguments[1])
+	assert.Equal(t, true, call.Arguments[2])
+	assert.Equal(t, 1, call.ReturnArguments[0])
+	assert.Equal(t, "two", call.ReturnArguments[1])
+	assert.Equal(t, true, call.ReturnArguments[2])
+	assert.Equal(t, 5, call.Repeatability)
+	assert.Nil(t, call.WaitFor)
 }
 
 func Test_Mock_Return_Nothing(t *testing.T) {
@@ -379,20 +454,19 @@ func Test_Mock_Return_Nothing(t *testing.T) {
 	// make a test impl object
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
-	assert.Equal(t, mockedService.On("TheExampleMethod", "A", "B", true).Return(), &mockedService.Mock)
+	c := mockedService.
+		On("TheExampleMethod", "A", "B", true).
+		Return()
 
-	// ensure the call was created
-	if assert.Equal(t, 1, len(mockedService.ExpectedCalls)) {
-		call := mockedService.ExpectedCalls[0]
+	require.Equal(t, []*Call{c}, mockedService.ExpectedCalls)
 
-		assert.Equal(t, "TheExampleMethod", call.Method)
-		assert.Equal(t, "A", call.Arguments[0])
-		assert.Equal(t, "B", call.Arguments[1])
-		assert.Equal(t, true, call.Arguments[2])
-		assert.Equal(t, 0, len(call.ReturnArguments))
-
-	}
+	call := mockedService.ExpectedCalls[0]
 
+	assert.Equal(t, "TheExampleMethod", call.Method)
+	assert.Equal(t, "A", call.Arguments[0])
+	assert.Equal(t, "B", call.Arguments[1])
+	assert.Equal(t, true, call.Arguments[2])
+	assert.Equal(t, 0, len(call.ReturnArguments))
 }
 
 func Test_Mock_findExpectedCall(t *testing.T) {
@@ -517,8 +591,13 @@ func Test_Mock_Called_For_Bounded_Repeatability(t *testing.T) {
 
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
-	mockedService.On("Test_Mock_Called_For_Bounded_Repeatability", 1, 2, 3).Return(5, "6", true).Once()
-	mockedService.On("Test_Mock_Called_For_Bounded_Repeatability", 1, 2, 3).Return(-1, "hi", false)
+	mockedService.
+		On("Test_Mock_Called_For_Bounded_Repeatability", 1, 2, 3).
+		Return(5, "6", true).
+		Once()
+	mockedService.
+		On("Test_Mock_Called_For_Bounded_Repeatability", 1, 2, 3).
+		Return(-1, "hi", false)
 
 	returnArguments1 := mockedService.Called(1, 2, 3)
 	returnArguments2 := mockedService.Called(1, 2, 3)
@@ -716,7 +795,9 @@ func Test_Mock_AssertCalled_WithAnythingOfTypeArgument(t *testing.T) {
 
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
 
-	mockedService.On("Test_Mock_AssertCalled_WithAnythingOfTypeArgument", Anything, Anything, Anything).Return()
+	mockedService.
+		On("Test_Mock_AssertCalled_WithAnythingOfTypeArgument", Anything, Anything, Anything).
+		Return()
 
 	mockedService.Called(1, "two", []uint8("three"))