/
testdialer.go
48 lines (40 loc) · 1.14 KB
/
testdialer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
/*
Copyright (c) Edgeless Systems GmbH
SPDX-License-Identifier: AGPL-3.0-only
*/
// Package testdialer provides a fake dialer for testing.
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, _, 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
}