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

package peerstream

import (
	"context"
	"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) {
	select {
	case err := <-c.ErrCh:
		return nil, err
	case r := <-c.ReplicationStream.sendCh:
		return r, nil
	case <-time.After(10 * time.Millisecond):
		return nil, io.EOF
	}
}

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
	mu  sync.Mutex
}

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

type incrementalTime struct {
	base time.Time
	next uint64
}

func (t *incrementalTime) Now() time.Time {
	t.next++
	return t.base.Add(time.Duration(t.next) * time.Second)
}