Mock can block returns with .After and .WaitUntil

You sometimes want to test concurrent behaviour like timeouts, but it
was not currently possible with our mocks.

This commits adds the following functions to Mock:

`After` will block m.Called for the given duration before returning.

```golang
func (m *Mock) After(d time.Duration) *Mock

Mock.On("MyMethod", arg1).After(1 * time.Millisecond)
```

`WaitUntil` is the primitive under `After`, and it's exposed in case you
want somebody wants more control over the returns. `Called` blocks until the w
channel is closed or receives a message.

```golang
func (m *Mock) WaitUntil(w <-chan time.Time) *Mock

Mock.On("MyMethod", arg1).WaitUntil(time.After(1 * time.Millisecond))
```

mock/mock.go

@@ -8,6 +8,7 @@ import (
 	"runtime"
 	"strings"
 	"sync"
+	"time"
 )
 
 // TestingT is an interface wrapper around *testing.T
@@ -37,6 +38,10 @@ type Call struct {
 	// The number of times to return the return arguments when setting
 	// expectations. 0 means to always return the value.
 	Repeatability int
+
+	// 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
 }
 
 // Mock is the workhorse used to track activity on another object.
@@ -95,7 +100,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})
+	m.ExpectedCalls = append(m.ExpectedCalls, Call{m.onMethodName, m.onMethodArguments, returnArguments, 0, nil})
 	return m
 }
 
@@ -121,6 +126,22 @@ func (m *Mock) Times(i int) {
 	m.ExpectedCalls[len(m.ExpectedCalls)-1].Repeatability = i
 }
 
+// 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.ExpectedCalls[len(m.ExpectedCalls)-1].WaitFor = w
+	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))
+}
+
 /*
 	Recording and responding to activity
 */
@@ -180,6 +201,7 @@ func callString(method string, arguments Arguments, includeArgumentValues bool)
 // Called tells the mock object that a method has been called, and gets an array
 // of arguments to return.  Panics if the call is unexpected (i.e. not preceeded by
 // appropriate .On .Return() calls)
+// If Call.WaitFor is set, blocks until the channel is closed or receives a message.
 func (m *Mock) Called(arguments ...interface{}) Arguments {
 	defer m.mutex.Unlock()
 	m.mutex.Lock()
@@ -220,7 +242,12 @@ func (m *Mock) Called(arguments ...interface{}) Arguments {
 	}
 
 	// add the call
-	m.Calls = append(m.Calls, Call{functionName, arguments, make([]interface{}, 0), 0})
+	m.Calls = append(m.Calls, Call{functionName, arguments, make([]interface{}, 0), 0, nil})
+
+	// block if specified
+	if call.WaitFor != nil {
+		<-call.WaitFor
+	}
 
 	return call.ReturnArguments
 

mock/mock_test.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"github.com/stretchr/testify/assert"
 	"testing"
+	"time"
 )
 
 /*
@@ -95,6 +96,58 @@ func Test_Mock_Return(t *testing.T) {
 		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) {
+
+	// make a test impl object
+	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)
+
+	// ensure the call was created
+	if assert.Equal(t, 1, len(mockedService.Mock.ExpectedCalls)) {
+		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.Equal(t, ch, call.WaitFor)
+
+	}
+
+}
+
+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)
+
+	// ensure the call was created
+	if assert.Equal(t, 1, len(mockedService.Mock.ExpectedCalls)) {
+		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)
 
 	}
 
@@ -119,6 +172,7 @@ func Test_Mock_Return_Once(t *testing.T) {
 		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)
 
 	}
 
@@ -143,6 +197,7 @@ func Test_Mock_Return_Twice(t *testing.T) {
 		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)
 
 	}
 
@@ -167,6 +222,7 @@ func Test_Mock_Return_Times(t *testing.T) {
 		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)
 
 	}
 
@@ -274,6 +330,43 @@ func Test_Mock_Called(t *testing.T) {
 
 }
 
+func asyncCall(m *Mock, ch chan Arguments) {
+	ch <- m.Called(1, 2, 3)
+}
+
+func Test_Mock_Called_blocks(t *testing.T) {
+
+	var mockedService *TestExampleImplementation = new(TestExampleImplementation)
+
+	mockedService.Mock.On("asyncCall", 1, 2, 3).Return(5, "6", true).After(2 * time.Millisecond)
+
+	ch := make(chan Arguments)
+
+	go asyncCall(&mockedService.Mock, ch)
+
+	select {
+	case <-ch:
+		t.Fatal("should have waited")
+	case <-time.After(1 * time.Millisecond):
+	}
+
+	returnArguments := <-ch
+
+	if assert.Equal(t, 1, len(mockedService.Mock.Calls)) {
+		assert.Equal(t, "asyncCall", mockedService.Mock.Calls[0].Method)
+		assert.Equal(t, 1, mockedService.Mock.Calls[0].Arguments[0])
+		assert.Equal(t, 2, mockedService.Mock.Calls[0].Arguments[1])
+		assert.Equal(t, 3, mockedService.Mock.Calls[0].Arguments[2])
+	}
+
+	if assert.Equal(t, 3, len(returnArguments)) {
+		assert.Equal(t, 5, returnArguments[0])
+		assert.Equal(t, "6", returnArguments[1])
+		assert.Equal(t, true, returnArguments[2])
+	}
+
+}
+
 func Test_Mock_Called_For_Bounded_Repeatability(t *testing.T) {
 
 	var mockedService *TestExampleImplementation = new(TestExampleImplementation)