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() }