open-consul/agent/grpc-external/services/peerstream/testing.go

147 lines
3.4 KiB
Go

package peerstream
import (
"context"
"fmt"
"io"
"sync"
"time"
"google.golang.org/grpc/metadata"
"github.com/hashicorp/consul/proto/pbpeerstream"
)
type MockClient struct {
mu sync.Mutex
ErrCh chan error
ReplicationStream *MockStream
}
func (c *MockClient) Send(r *pbpeerstream.ReplicationMessage) error {
c.ReplicationStream.recvCh <- r
return nil
}
func (c *MockClient) Recv() (*pbpeerstream.ReplicationMessage, error) {
return c.RecvWithTimeout(10 * time.Millisecond)
}
func (c *MockClient) RecvWithTimeout(dur time.Duration) (*pbpeerstream.ReplicationMessage, error) {
select {
case err := <-c.ErrCh:
return nil, err
case r := <-c.ReplicationStream.sendCh:
return r, nil
case <-time.After(dur):
return nil, io.EOF
}
}
func (c *MockClient) Close() {
close(c.ReplicationStream.recvCh)
}
func NewMockClient(ctx context.Context) *MockClient {
return &MockClient{
ReplicationStream: newTestReplicationStream(ctx),
}
}
// MockStream mocks peering.PeeringService_StreamResourcesServer
type MockStream struct {
sendCh chan *pbpeerstream.ReplicationMessage
recvCh chan *pbpeerstream.ReplicationMessage
ctx context.Context
}
var _ pbpeerstream.PeerStreamService_StreamResourcesServer = (*MockStream)(nil)
func newTestReplicationStream(ctx context.Context) *MockStream {
return &MockStream{
sendCh: make(chan *pbpeerstream.ReplicationMessage, 1),
recvCh: make(chan *pbpeerstream.ReplicationMessage, 1),
ctx: ctx,
}
}
// Send implements pbpeerstream.PeeringService_StreamResourcesServer
func (s *MockStream) Send(r *pbpeerstream.ReplicationMessage) error {
s.sendCh <- r
return nil
}
// Recv implements pbpeerstream.PeeringService_StreamResourcesServer
func (s *MockStream) Recv() (*pbpeerstream.ReplicationMessage, error) {
r := <-s.recvCh
if r == nil {
return nil, io.EOF
}
return r, nil
}
// Context implements grpc.ServerStream and grpc.ClientStream
func (s *MockStream) Context() context.Context {
return s.ctx
}
// SendMsg implements grpc.ServerStream and grpc.ClientStream
func (s *MockStream) SendMsg(m interface{}) error {
return nil
}
// RecvMsg implements grpc.ServerStream and grpc.ClientStream
func (s *MockStream) RecvMsg(m interface{}) error {
return nil
}
// SetHeader implements grpc.ServerStream
func (s *MockStream) SetHeader(metadata.MD) error {
return nil
}
// SendHeader implements grpc.ServerStream
func (s *MockStream) SendHeader(metadata.MD) error {
return nil
}
// SetTrailer implements grpc.ServerStream
func (s *MockStream) SetTrailer(metadata.MD) {}
// incrementalTime is an artificial clock used during testing. For those
// scenarios you would pass around the method pointer for `Now` in places where
// you would be using `time.Now`.
type incrementalTime struct {
base time.Time
next uint64
mu sync.Mutex
}
// Now advances the internal clock by 1 second and returns that value.
func (t *incrementalTime) Now() time.Time {
t.mu.Lock()
defer t.mu.Unlock()
t.next++
dur := time.Duration(t.next) * time.Second
return t.base.Add(dur)
}
// FutureNow will return a given future value of the Now() function.
// The numerical argument indicates which future Now value you wanted. The
// value must be > 0.
func (t *incrementalTime) FutureNow(n int) time.Time {
if n < 1 {
panic(fmt.Sprintf("argument must be > 1 but was %d", n))
}
t.mu.Lock()
defer t.mu.Unlock()
dur := time.Duration(t.next+uint64(n)) * time.Second
return t.base.Add(dur)
}