constellation/coordinator/util/testdialer/bufconndialer.go

42 lines
1.0 KiB
Go
Raw Normal View History

package testdialer
import (
"context"
"fmt"
"net"
"sync"
"google.golang.org/grpc/test/bufconn"
)
// BufconnDialer is a fake dialer based on gRPC bufconn package.
type BufconnDialer struct {
mut sync.Mutex
listeners map[string]*bufconn.Listener
}
// NewBufconnDialer creates a new bufconn dialer for testing.
func NewBufconnDialer() *BufconnDialer {
return &BufconnDialer{listeners: make(map[string]*bufconn.Listener)}
}
// DialContext implements the Dialer interface.
func (b *BufconnDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
b.mut.Lock()
listener, ok := b.listeners[address]
b.mut.Unlock()
if !ok {
return nil, fmt.Errorf("could not connect to server on %v", address)
}
return listener.DialContext(ctx)
}
// GetListener returns a fake listener that is coupled with this dialer.
func (b *BufconnDialer) GetListener(endpoint string) net.Listener {
listener := bufconn.Listen(1024)
b.mut.Lock()
b.listeners[endpoint] = listener
b.mut.Unlock()
return listener
}