Merge pull request #148 from ernesto-jimenez/master

Mock can block returns with .After and .WaitUntil

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)