constellation/coordinator/util/testdialer/bufconndialer.go
Leonard Cohnen 2d8fcd9bf4 monorepo
Co-authored-by: Malte Poll <mp@edgeless.systems>
Co-authored-by: katexochen <katexochen@users.noreply.github.com>
Co-authored-by: Daniel Weiße <dw@edgeless.systems>
Co-authored-by: Thomas Tendyck <tt@edgeless.systems>
Co-authored-by: Benedict Schlueter <bs@edgeless.systems>
Co-authored-by: leongross <leon.gross@rub.de>
Co-authored-by: Moritz Eckert <m1gh7ym0@gmail.com>
2022-03-22 16:09:39 +01:00

42 lines
1.0 KiB
Go

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
}