/
bufconn.go
96 lines (80 loc) · 2.81 KB
/
bufconn.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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
package mock
import (
"context"
"net"
"google.golang.org/grpc"
"google.golang.org/grpc/test/bufconn"
)
const (
bufsize = 1024 * 1024
buftarget = "bufnet"
)
// Listener handles gRPC connections using an in-memory buffer that is useful for
// testing to prevent actual TCP network requests. Using a bufconn connection provides
// the most realistic gRPC server for tests that include serialization and
// deserialization of protocol buffers and an actual wire transfer between client and
// server. We prefer to use the bufconn over simply making method calls to the handlers.
type Listener struct {
sock *bufconn.Listener
target string
}
// New creates a bufconn listener ready to attach servers and clients to. To provide a
// different target name (e.g. for mTLS buffers) use the WithTarget() dial option. You
// can also specify a different buffer size using WithBufferSize() or pass in an already
// instantiated bufconn.Listener using WithBuffer().
func NewBufConn(opts ...DialOption) *Listener {
sock := &Listener{}
for _, opt := range opts {
opt(sock)
}
if sock.target == "" {
sock.target = buftarget
}
if sock.sock == nil {
sock.sock = bufconn.Listen(bufsize)
}
return sock
}
// Sock returns the server side of the bufconn connection.
func (l *Listener) Sock() net.Listener {
return l.sock
}
// Close the bufconn listener and prevent either clients or servers from communicating.
func (l *Listener) Close() error {
return l.sock.Close()
}
// Connect returns the client side of the bufconn connection.
func (l *Listener) Connect(ctx context.Context, opts ...grpc.DialOption) (cc *grpc.ClientConn, err error) {
opts = append([]grpc.DialOption{grpc.WithContextDialer(l.Dialer)}, opts...)
if cc, err = grpc.DialContext(ctx, l.target, opts...); err != nil {
return nil, err
}
return cc, nil
}
// Dialer implements the ContextDialer interface for use with grpc.DialOptions
func (l *Listener) Dialer(context.Context, string) (net.Conn, error) {
return l.sock.Dial()
}
// DialOption -- optional arguments for constructing a Listener.
type DialOption func(*Listener)
// WithTarget allows the user to change the "endpoint" of the connection from "bufnet"
// to some other endpoint. This is useful for tests that include mTLS to ensure that
// TLS certificate handling is correct.
func WithTarget(target string) DialOption {
return func(lis *Listener) {
lis.target = target
}
}
// The default buffer size is 1MiB -- if ou need a larger buffer to send larger messages
// then specify this dial option with a larger size.
func WithBufferSize(size int) DialOption {
return func(lis *Listener) {
lis.sock = bufconn.Listen(size)
}
}
// Allows you to pass an already instantiated grpc bufconn.Listener into the Listener.
func WithBuffer(sock *bufconn.Listener) DialOption {
return func(lis *Listener) {
lis.sock = sock
}
}