Fixing After(duration) to wait during method call

mock/mock.go

@@ -56,6 +56,8 @@ type Call struct {
 	// receives a message or is closed. nil means it returns immediately.
 	WaitFor <-chan time.Time
 
+	waitTime time.Duration
+
 	// Holds a handler used to manipulate arguments content that are passed by
 	// reference. It's useful when mocking methods such as unmarshalers or
 	// decoders.
@@ -134,7 +136,10 @@ func (c *Call) WaitUntil(w <-chan time.Time) *Call {
 //
 //    Mock.On("MyMethod", arg1, arg2).After(time.Second)
 func (c *Call) After(d time.Duration) *Call {
-	return c.WaitUntil(time.After(d))
+	c.lock()
+	defer c.unlock()
+	c.waitTime = d
+	return c
 }
 
 // Run sets a handler to be called before returning. It can be used when
@@ -327,18 +332,12 @@ func (m *Mock) MethodCalled(methodName string, arguments ...interface{}) Argumen
 		}
 	}
 
-	switch {
-	case call.Repeatability == 1:
+	if call.Repeatability == 1 {
 		call.Repeatability = -1
-		call.totalCalls++
-
-	case call.Repeatability > 1:
+	} else if call.Repeatability > 1 {
 		call.Repeatability--
-		call.totalCalls++
-
-	case call.Repeatability == 0:
-		call.totalCalls++
 	}
+	call.totalCalls++
 
 	// add the call
 	m.Calls = append(m.Calls, *newCall(m, methodName, arguments...))
@@ -347,6 +346,8 @@ func (m *Mock) MethodCalled(methodName string, arguments ...interface{}) Argumen
 	// block if specified
 	if call.WaitFor != nil {
 		<-call.WaitFor
+	} else {
+		time.Sleep(call.waitTime)
 	}
 
 	m.mutex.Lock()

mock/mock_test.go

@@ -2,6 +2,7 @@ package mock
 
 import (
 	"errors"
+	"fmt"
 	"sync"
 	"testing"
 	"time"
@@ -1316,6 +1317,36 @@ func Test_MockReturnAndCalledConcurrent(t *testing.T) {
 	wg.Wait()
 }
 
+type timer struct{ Mock }
+
+func (s *timer) GetTime(i int) string {
+	return s.Called(i).Get(0).(string)
+}
+
+func TestAfterTotalWaitTimeWhileExecution(t *testing.T) {
+	waitDuration := 1
+	total, waitMs := 5, time.Millisecond*time.Duration(waitDuration)
+	aTimer := new(timer)
+	for i := 0; i < total; i++ {
+		aTimer.On("GetTime", i).After(waitMs).Return(fmt.Sprintf("Time%d", i)).Once()
+	}
+	time.Sleep(waitMs)
+	start := time.Now()
+	var results []string
+
+	for i := 0; i < total; i++ {
+		results = append(results, aTimer.GetTime(i))
+	}
+
+	end := time.Now()
+	elapsedTime := end.Sub(start)
+	assert.True(t, elapsedTime > waitMs, fmt.Sprintf("Total elapsed time:%v should be atleast greater than %v", elapsedTime, waitMs))
+	assert.Equal(t, total, len(results))
+	for i, _ := range results {
+		assert.Equal(t, fmt.Sprintf("Time%d", i), results[i], "Return value of method should be same")
+	}
+}
+
 func ConcurrencyTestMethod(m *Mock) {
 	m.Called()
 }