Use go.nhat.io/wait (#59)

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/swaggest/assertjson v1.7.0
 	go.nhat.io/aferomock v0.4.0
 	go.nhat.io/matcher/v2 v2.0.0
+	go.nhat.io/wait v0.1.0
 	google.golang.org/grpc v1.52.3
 	google.golang.org/protobuf v1.28.1
 )
@@ -19,6 +20,7 @@ require (
 	github.com/fsnotify/fsnotify v1.4.9 // indirect
 	github.com/golang/protobuf v1.5.2 // indirect
 	github.com/iancoleman/orderedmap v0.2.0 // indirect
+	github.com/kr/pretty v0.2.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/sergi/go-diff v1.3.1 // indirect
 	github.com/stretchr/objx v0.5.0 // indirect

go.sum

@@ -539,6 +539,8 @@ github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxv
 github.com/kr/fs v0.1.0/go.mod h1:FFnZGqtBN9Gxj7eW1uZ42v5BccTP0vu6NEaFoC2HwRg=
 github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
+github.com/kr/pretty v0.2.1 h1:Fmg33tUaq4/8ym9TJN1x7sLJnHVwhP33CNkpYV/7rwI=
+github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
@@ -548,7 +550,6 @@ github.com/mattn/go-colorable v0.1.8/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope
 github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU=
 github.com/mattn/go-isatty v0.0.14 h1:yVuAays6BHfxijgZPzw+3Zlu5yQgKGP2/hcQbHb7S9Y=
 github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94=
-github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e h1:fD57ERR4JtEqsWbfPhv4DMiApHyliiK5xCTNVSPiaAs=
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
 github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
 github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE=
@@ -612,6 +613,8 @@ go.nhat.io/aferomock v0.4.0 h1:gs3nJzIqAezglUuaPfautAmZwulwRWLcfSSzdK4YCC0=
 go.nhat.io/aferomock v0.4.0/go.mod h1:msi5MDOtJ/AroUa/lDc3jVGOILM4SKP//4yBRImOvkI=
 go.nhat.io/matcher/v2 v2.0.0 h1:W+rbHi0hKuZHtOQH4U5g+KwyKyfVioIxrxjoGRcUETE=
 go.nhat.io/matcher/v2 v2.0.0/go.mod h1:cL5oYp0M9A4L8jEGqjmUfy+k7AXVDddoVt6aYIL1r5g=
+go.nhat.io/wait v0.1.0 h1:aQ4YDzaOgFbypiJ9c/eAfOIB1G25VOv7Gd2QS8uz1gw=
+go.nhat.io/wait v0.1.0/go.mod h1:+ijMghc9/9zXi+HDcs49HNReprvXOZha2Q3jTOtqJrE=
 go.opencensus.io v0.21.0/go.mod h1:mSImk1erAIZhrmZN+AvHh14ztQfjbGwt4TtuofqLduU=
 go.opencensus.io v0.22.0/go.mod h1:+kGneAE2xo2IficOXnaByMWTGM9T73dGwxeWcUqIpI8=
 go.opencensus.io v0.22.2/go.mod h1:yxeiOL68Rb0Xd1ddK5vPZ/oVn4vY4Ynel7k9FzqtOIw=
@@ -1153,8 +1156,9 @@ google.golang.org/protobuf v1.28.1/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqw
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f h1:BLraFXnmrev5lT+xlilqcH8XK9/i0At2xKjWk4p6zsU=
 gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
 gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
 gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys=
 gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ=

server_expectation.go

@@ -1,11 +1,10 @@
 package grpcmock
 
 import (
-	"context"
 	"sync"
-	"time"
 
 	"github.com/spf13/afero"
+	"go.nhat.io/wait"
 	"google.golang.org/grpc/codes"
 
 	xmatcher "go.nhat.io/grpcmock/matcher"
@@ -15,7 +14,7 @@ import (
 type baseExpectation struct {
 	locker sync.Locker
 	fs     afero.Fs
-	delay  delayer
+	waiter wait.Waiter
 
 	serviceDesc *service.Method
 
@@ -95,35 +94,3 @@ func (e *baseExpectation) FulfilledTimes() uint {
 
 	return e.fulfilledTimes
 }
-
-type delayer func(ctx context.Context) error
-
-func noWait() delayer {
-	return func(ctx context.Context) error {
-		return ctx.Err()
-	}
-}
-
-func waitForSignal(c <-chan time.Time) delayer {
-	return func(ctx context.Context) error {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-
-		case <-c:
-			return nil
-		}
-	}
-}
-
-func waitForDuration(d time.Duration) delayer {
-	return func(ctx context.Context) error {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-
-		case <-time.After(d):
-			return nil
-		}
-	}
-}

server_expectation_bidirectional_stream.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/spf13/afero"
 	"go.nhat.io/matcher/v2"
+	"go.nhat.io/wait"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -207,7 +208,7 @@ func (e *bidirectionalStreamExpectation) Run(handler func(ctx context.Context, s
 }
 
 func (e *bidirectionalStreamExpectation) Handle(ctx context.Context, in interface{}, _ interface{}) error {
-	if err := e.delay(ctx); err != nil {
+	if err := e.waiter.Wait(ctx); err != nil {
 		return xerrors.StatusError(err)
 	}
 
@@ -285,7 +286,7 @@ func (e *bidirectionalStreamExpectation) WaitUntil(w <-chan time.Time) Bidirecti
 	e.lock()
 	defer e.unlock()
 
-	e.delay = waitForSignal(w)
+	e.waiter = wait.ForSignal(w)
 
 	return e
 }
@@ -301,7 +302,7 @@ func (e *bidirectionalStreamExpectation) After(d time.Duration) BidirectionalStr
 	e.lock()
 	defer e.unlock()
 
-	e.delay = waitForDuration(d)
+	e.waiter = wait.ForDuration(d)
 
 	return e
 }
@@ -312,7 +313,7 @@ func newBidirectionalStreamExpectation(svc *service.Method) *bidirectionalStream
 		baseExpectation: &baseExpectation{
 			locker:      &sync.Mutex{},
 			fs:          afero.NewOsFs(),
-			delay:       noWait(),
+			waiter:      wait.NoWait,
 			serviceDesc: svc,
 		},
 		run: func(context.Context, grpc.ServerStream) error {

server_expectation_client_stream.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/spf13/afero"
 	"go.nhat.io/matcher/v2"
+	"go.nhat.io/wait"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -290,7 +291,7 @@ func (e *clientStreamExpectation) Run(handler func(ctx context.Context, s grpc.S
 
 // Handle executes the GRPC request.
 func (e *clientStreamExpectation) Handle(ctx context.Context, in interface{}, out interface{}) error {
-	if err := e.delay(ctx); err != nil {
+	if err := e.waiter.Wait(ctx); err != nil {
 		return xerrors.StatusError(err)
 	}
 
@@ -345,7 +346,7 @@ func (e *clientStreamExpectation) WaitUntil(w <-chan time.Time) ClientStreamExpe
 	e.lock()
 	defer e.unlock()
 
-	e.delay = waitForSignal(w)
+	e.waiter = wait.ForSignal(w)
 
 	return e
 }
@@ -354,7 +355,7 @@ func (e *clientStreamExpectation) After(d time.Duration) ClientStreamExpectation
 	e.lock()
 	defer e.unlock()
 
-	e.delay = waitForDuration(d)
+	e.waiter = wait.ForDuration(d)
 
 	return e
 }
@@ -364,7 +365,7 @@ func newClientStreamExpectation(svc *service.Method) *clientStreamExpectation {
 	return &clientStreamExpectation{
 		baseExpectation: &baseExpectation{
 			locker:      &sync.Mutex{},
-			delay:       noWait(),
+			waiter:      wait.NoWait,
 			fs:          afero.NewOsFs(),
 			serviceDesc: svc,
 		},

server_expectation_server_stream.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/spf13/afero"
 	"go.nhat.io/matcher/v2"
+	"go.nhat.io/wait"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
@@ -389,7 +390,7 @@ func (e *serverStreamExpectation) Run(handler func(ctx context.Context, in inter
 }
 
 func (e *serverStreamExpectation) Handle(ctx context.Context, in interface{}, out interface{}) error {
-	if err := e.delay(ctx); err != nil {
+	if err := e.waiter.Wait(ctx); err != nil {
 		return xerrors.StatusError(err)
 	}
 
@@ -422,7 +423,7 @@ func (e *serverStreamExpectation) WaitUntil(w <-chan time.Time) ServerStreamExpe
 	e.lock()
 	defer e.unlock()
 
-	e.delay = waitForSignal(w)
+	e.waiter = wait.ForSignal(w)
 
 	return e
 }
@@ -431,7 +432,7 @@ func (e *serverStreamExpectation) After(d time.Duration) ServerStreamExpectation
 	e.lock()
 	defer e.unlock()
 
-	e.delay = waitForDuration(d)
+	e.waiter = wait.ForDuration(d)
 
 	return e
 }
@@ -442,7 +443,7 @@ func newServerStreamExpectation(svc *service.Method) *serverStreamExpectation {
 		baseExpectation: &baseExpectation{
 			locker:      &sync.Mutex{},
 			fs:          afero.NewOsFs(),
-			delay:       noWait(),
+			waiter:      wait.NoWait,
 			serviceDesc: svc,
 		},
 		run: func(context.Context, interface{}, grpc.ServerStream) error {
@@ -663,6 +664,6 @@ func stepReturnErrorf(code codes.Code, msg string, args ...interface{}) serverSt
 
 func stepWait(d time.Duration) serverStreamHandlerStepFunc {
 	return func(ctx context.Context, _ grpc.ServerStream) error {
-		return waitForDuration(d)(ctx)
+		return wait.ForDuration(d).Wait(ctx)
 	}
 }

server_expectation_unary.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/spf13/afero"
 	"go.nhat.io/matcher/v2"
+	"go.nhat.io/wait"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/encoding/protojson"
@@ -315,7 +316,7 @@ func (e *unaryExpectation) Run(handler func(ctx context.Context, in interface{})
 }
 
 func (e *unaryExpectation) Handle(ctx context.Context, in interface{}, out interface{}) error {
-	if err := e.delay(ctx); err != nil {
+	if err := e.waiter.Wait(ctx); err != nil {
 		return xerrors.StatusError(err)
 	}
 
@@ -368,7 +369,7 @@ func (e *unaryExpectation) WaitUntil(w <-chan time.Time) UnaryExpectation {
 	e.lock()
 	defer e.unlock()
 
-	e.delay = waitForSignal(w)
+	e.waiter = wait.ForSignal(w)
 
 	return e
 }
@@ -377,7 +378,7 @@ func (e *unaryExpectation) After(d time.Duration) UnaryExpectation {
 	e.lock()
 	defer e.unlock()
 
-	e.delay = waitForDuration(d)
+	e.waiter = wait.ForDuration(d)
 
 	return e
 }
@@ -388,7 +389,7 @@ func newUnaryExpectation(svc *service.Method) *unaryExpectation {
 		baseExpectation: &baseExpectation{
 			locker:      &sync.Mutex{},
 			fs:          afero.NewOsFs(),
-			delay:       noWait(),
+			waiter:      wait.NoWait,
 			serviceDesc: svc,
 		},
 		run: func(ctx context.Context, in interface{}) (interface{}, error) {