Merge pull request #76 from bytedance/fix/race

fix: race on sequence

mock_sequence.go

@@ -16,7 +16,11 @@
 
 package mockey
 
-import "github.com/bytedance/mockey/internal/tool"
+import (
+	"sync"
+
+	"github.com/bytedance/mockey/internal/tool"
+)
 
 type SequenceOpt interface {
 	// Private make sure it is mockey private interface
@@ -36,6 +40,7 @@ type sequence struct {
 	values        []*sequenceValue
 	curV          int // current value
 	curT          int // current value times
+	readLock      sync.Mutex
 }
 
 type sequenceValue struct {
@@ -44,6 +49,7 @@ type sequenceValue struct {
 }
 
 func (s *sequence) GetNext() []interface{} {
+	s.readLock.Lock()
 	seqV := s.values[s.curV]
 	s.curT++
 	if s.curT >= seqV.t {
@@ -53,6 +59,7 @@ func (s *sequence) GetNext() []interface{} {
 			s.curV = 0
 		}
 	}
+	s.readLock.Unlock()
 	return seqV.v
 }
 

mock_sequence_test.go

@@ -18,11 +18,34 @@ package mockey
 
 import (
 	"fmt"
+	"sync"
 	"testing"
 
 	"github.com/smartystreets/goconvey/convey"
 )
 
+func TestSequenceRace(t *testing.T) {
+	PatchConvey("test sequence race", t, func(c convey.C) {
+		fn := func() int { return -1 }
+		mocker := Mock(fn).Return(Sequence(0).Then(1)).Build()
+
+		wg := sync.WaitGroup{}
+		wg.Add(2)
+		go func() {
+			defer wg.Done()
+			v := fn()
+			c.So(v, convey.ShouldNotEqual, -1)
+		}()
+		go func() {
+			defer wg.Done()
+			v := fn()
+			c.So(v, convey.ShouldNotEqual, -1)
+		}()
+		wg.Wait()
+		c.So(mocker.MockTimes(), convey.ShouldEqual, 2)
+	})
+}
+
 func TestSequenceOpt(t *testing.T) {
 	PatchConvey("test sequenceOpt", t, func() {
 		fn := func() (string, int) {