From b3596e92796169eee4a0050bceaa767a48cb0075 Mon Sep 17 00:00:00 2001
From: Dinesh Kumar <dinesh.kumar@go-jek.com>
Date: Sat, 30 Dec 2017 21:00:04 +0530
Subject: [PATCH] Fxinng After(time.Duration) to wait properly

commit 5b0291d47dc3a70cc6c6be3bba3c2f934a8e933e
Merge: 1f324ec 8ccf48a
Author: Dinesh Kumar <dineshkumar-cse@users.noreply.github.com>
Date:   Sat Dec 30 19:55:13 2017 +0530

    Merge branch 'master' into master

commit 1f324ec8cbe892a2c5364904643aa1b710121824
Author: Dinesh Kumar <dinesh.kumar@go-jek.com>
Date:   Sat Dec 30 19:07:16 2017 +0530

    Fixing comments: reduced test time, locking in after, unexported waitTime

    - WaitUntil/After overrides each other value
    - currently if channel is set that takes the priority, if not the
      waitTime is used to Sleep

commit a7101ec14224918fc06532acae76418d1b4f04c5
Author: Dinesh Kumar <dinesh.kumar@go-jek.com>
Date:   Fri Dec 29 13:02:14 2017 +0530

    Using if else instead of switch and pulling out the commone one

commit 936f63dd689acacfd051867c5c700489caa1dbc2
Author: Dinesh Kumar <dinesh.kumar@go-jek.com>
Date:   Fri Dec 29 12:41:40 2017 +0530

    After - making it wait during method call

    After was using call.WaitUntil(time.After(duration)),
    <-time.After(duration) returns a channel with the timer immediately
    started
---
 mock/mock.go      | 21 +++++++++++----------
 mock/mock_test.go | 31 +++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/mock/mock.go b/mock/mock.go
index 08e41e8..208b838 100644
--- a/mock/mock.go
+++ b/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()
diff --git a/mock/mock_test.go b/mock/mock_test.go
index 0070f94..cb245ba 100644
--- a/mock/mock_test.go
+++ b/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()
 }