forked from hashicorp/go-plugin
-
Notifications
You must be signed in to change notification settings - Fork 0
/
testing.go
183 lines (154 loc) · 4.59 KB
/
testing.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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0
package plugin
import (
"bytes"
"context"
"io"
"net"
"net/rpc"
hclog "github.com/hashicorp/go-hclog"
"github.com/hashicorp/go-plugin/internal/plugin"
"github.com/mitchellh/go-testing-interface"
"google.golang.org/grpc"
)
// TestOptions allows specifying options that can affect the behavior of the
// test functions
type TestOptions struct {
//ServerStdout causes the given value to be used in place of a blank buffer
//for RPCServer's Stdout
ServerStdout io.ReadCloser
//ServerStderr causes the given value to be used in place of a blank buffer
//for RPCServer's Stderr
ServerStderr io.ReadCloser
}
// The testing file contains test helpers that you can use outside of
// this package for making it easier to test plugins themselves.
// TestConn is a helper function for returning a client and server
// net.Conn connected to each other.
func TestConn(t testing.T) (net.Conn, net.Conn) {
// Listen to any local port. This listener will be closed
// after a single connection is established.
l, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("err: %s", err)
}
// Start a goroutine to accept our client connection
var serverConn net.Conn
doneCh := make(chan struct{})
go func() {
defer close(doneCh)
defer l.Close()
var err error
serverConn, err = l.Accept()
if err != nil {
t.Fatalf("err: %s", err)
}
}()
// Connect to the server
clientConn, err := net.Dial("tcp", l.Addr().String())
if err != nil {
t.Fatalf("err: %s", err)
}
// Wait for the server side to acknowledge it has connected
<-doneCh
return clientConn, serverConn
}
// TestRPCConn returns a rpc client and server connected to each other.
func TestRPCConn(t testing.T) (*rpc.Client, *rpc.Server) {
clientConn, serverConn := TestConn(t)
server := rpc.NewServer()
go server.ServeConn(serverConn)
client := rpc.NewClient(clientConn)
return client, server
}
// TestPluginRPCConn returns a plugin RPC client and server that are connected
// together and configured.
func TestPluginRPCConn(t testing.T, ps map[string]Plugin, opts *TestOptions) (*RPCClient, *RPCServer) {
// Create two net.Conns we can use to shuttle our control connection
clientConn, serverConn := TestConn(t)
// Start up the server
server := &RPCServer{Plugins: ps, Stdout: new(bytes.Buffer), Stderr: new(bytes.Buffer)}
if opts != nil {
if opts.ServerStdout != nil {
server.Stdout = opts.ServerStdout
}
if opts.ServerStderr != nil {
server.Stderr = opts.ServerStderr
}
}
go server.ServeConn(serverConn)
// Connect the client to the server
client, err := NewRPCClient(clientConn, ps)
if err != nil {
t.Fatalf("err: %s", err)
}
return client, server
}
// TestGRPCConn returns a gRPC client conn and grpc server that are connected
// together and configured. The register function is used to register services
// prior to the Serve call. This is used to test gRPC connections.
func TestGRPCConn(t testing.T, register func(*grpc.Server)) (*grpc.ClientConn, *grpc.Server) {
// Create a listener
l, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("err: %s", err)
}
server := grpc.NewServer()
register(server)
go server.Serve(l)
// Connect to the server
conn, err := grpc.Dial(
l.Addr().String(),
grpc.WithBlock(),
grpc.WithInsecure())
if err != nil {
t.Fatalf("err: %s", err)
}
// Connection successful, close the listener
l.Close()
return conn, server
}
// TestPluginGRPCConn returns a plugin gRPC client and server that are connected
// together and configured. This is used to test gRPC connections.
func TestPluginGRPCConn(t testing.T, ps map[string]Plugin) (*GRPCClient, *GRPCServer) {
// Create a listener
l, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("err: %s", err)
}
// Start up the server
server := &GRPCServer{
Plugins: ps,
DoneCh: make(chan struct{}),
Server: DefaultGRPCServer,
Stdout: new(bytes.Buffer),
Stderr: new(bytes.Buffer),
logger: hclog.Default(),
}
if err := server.Init(); err != nil {
t.Fatalf("err: %s", err)
}
go server.Serve(l)
// Connect to the server
conn, err := grpc.Dial(
l.Addr().String(),
grpc.WithBlock(),
grpc.WithInsecure())
if err != nil {
t.Fatalf("err: %s", err)
}
brokerGRPCClient := newGRPCBrokerClient(conn)
broker := newGRPCBroker(brokerGRPCClient, nil)
go broker.Run()
go brokerGRPCClient.StartStream()
// Create the client
client := &GRPCClient{
Conn: conn,
Plugins: ps,
broker: broker,
doneCtx: context.Background(),
controller: plugin.NewGRPCControllerClient(conn),
}
return client, server
}