mirror of https://github.com/hashicorp/consul
129 lines
2.8 KiB
Go
129 lines
2.8 KiB
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)
|
|
}
|