]>
Commit | Line | Data |
---|---|---|
bae9f6d2 JC |
1 | package plugin |
2 | ||
3 | import ( | |
4 | "bytes" | |
15c0b25d | 5 | "context" |
107c1cdb | 6 | "io" |
bae9f6d2 JC |
7 | "net" |
8 | "net/rpc" | |
15c0b25d AP |
9 | |
10 | "github.com/mitchellh/go-testing-interface" | |
107c1cdb ND |
11 | hclog "github.com/hashicorp/go-hclog" |
12 | "github.com/hashicorp/go-plugin/internal/plugin" | |
15c0b25d | 13 | "google.golang.org/grpc" |
bae9f6d2 JC |
14 | ) |
15 | ||
107c1cdb ND |
16 | // TestOptions allows specifying options that can affect the behavior of the |
17 | // test functions | |
18 | type TestOptions struct { | |
19 | //ServerStdout causes the given value to be used in place of a blank buffer | |
20 | //for RPCServer's Stdout | |
21 | ServerStdout io.ReadCloser | |
22 | ||
23 | //ServerStderr causes the given value to be used in place of a blank buffer | |
24 | //for RPCServer's Stderr | |
25 | ServerStderr io.ReadCloser | |
26 | } | |
27 | ||
bae9f6d2 JC |
28 | // The testing file contains test helpers that you can use outside of |
29 | // this package for making it easier to test plugins themselves. | |
30 | ||
31 | // TestConn is a helper function for returning a client and server | |
32 | // net.Conn connected to each other. | |
15c0b25d | 33 | func TestConn(t testing.T) (net.Conn, net.Conn) { |
bae9f6d2 JC |
34 | // Listen to any local port. This listener will be closed |
35 | // after a single connection is established. | |
36 | l, err := net.Listen("tcp", "127.0.0.1:0") | |
37 | if err != nil { | |
38 | t.Fatalf("err: %s", err) | |
39 | } | |
40 | ||
41 | // Start a goroutine to accept our client connection | |
42 | var serverConn net.Conn | |
43 | doneCh := make(chan struct{}) | |
44 | go func() { | |
45 | defer close(doneCh) | |
46 | defer l.Close() | |
47 | var err error | |
48 | serverConn, err = l.Accept() | |
49 | if err != nil { | |
50 | t.Fatalf("err: %s", err) | |
51 | } | |
52 | }() | |
53 | ||
54 | // Connect to the server | |
55 | clientConn, err := net.Dial("tcp", l.Addr().String()) | |
56 | if err != nil { | |
57 | t.Fatalf("err: %s", err) | |
58 | } | |
59 | ||
60 | // Wait for the server side to acknowledge it has connected | |
61 | <-doneCh | |
62 | ||
63 | return clientConn, serverConn | |
64 | } | |
65 | ||
66 | // TestRPCConn returns a rpc client and server connected to each other. | |
15c0b25d | 67 | func TestRPCConn(t testing.T) (*rpc.Client, *rpc.Server) { |
bae9f6d2 JC |
68 | clientConn, serverConn := TestConn(t) |
69 | ||
70 | server := rpc.NewServer() | |
71 | go server.ServeConn(serverConn) | |
72 | ||
73 | client := rpc.NewClient(clientConn) | |
74 | return client, server | |
75 | } | |
76 | ||
77 | // TestPluginRPCConn returns a plugin RPC client and server that are connected | |
78 | // together and configured. | |
107c1cdb | 79 | func TestPluginRPCConn(t testing.T, ps map[string]Plugin, opts *TestOptions) (*RPCClient, *RPCServer) { |
bae9f6d2 JC |
80 | // Create two net.Conns we can use to shuttle our control connection |
81 | clientConn, serverConn := TestConn(t) | |
82 | ||
83 | // Start up the server | |
84 | server := &RPCServer{Plugins: ps, Stdout: new(bytes.Buffer), Stderr: new(bytes.Buffer)} | |
107c1cdb ND |
85 | if opts != nil { |
86 | if opts.ServerStdout != nil { | |
87 | server.Stdout = opts.ServerStdout | |
88 | } | |
89 | if opts.ServerStderr != nil { | |
90 | server.Stderr = opts.ServerStderr | |
91 | } | |
92 | } | |
bae9f6d2 JC |
93 | go server.ServeConn(serverConn) |
94 | ||
95 | // Connect the client to the server | |
96 | client, err := NewRPCClient(clientConn, ps) | |
97 | if err != nil { | |
98 | t.Fatalf("err: %s", err) | |
99 | } | |
100 | ||
101 | return client, server | |
102 | } | |
15c0b25d AP |
103 | |
104 | // TestGRPCConn returns a gRPC client conn and grpc server that are connected | |
105 | // together and configured. The register function is used to register services | |
106 | // prior to the Serve call. This is used to test gRPC connections. | |
107 | func TestGRPCConn(t testing.T, register func(*grpc.Server)) (*grpc.ClientConn, *grpc.Server) { | |
108 | // Create a listener | |
109 | l, err := net.Listen("tcp", "127.0.0.1:0") | |
110 | if err != nil { | |
111 | t.Fatalf("err: %s", err) | |
112 | } | |
113 | ||
114 | server := grpc.NewServer() | |
115 | register(server) | |
116 | go server.Serve(l) | |
117 | ||
118 | // Connect to the server | |
119 | conn, err := grpc.Dial( | |
120 | l.Addr().String(), | |
121 | grpc.WithBlock(), | |
122 | grpc.WithInsecure()) | |
123 | if err != nil { | |
124 | t.Fatalf("err: %s", err) | |
125 | } | |
126 | ||
127 | // Connection successful, close the listener | |
128 | l.Close() | |
129 | ||
130 | return conn, server | |
131 | } | |
132 | ||
133 | // TestPluginGRPCConn returns a plugin gRPC client and server that are connected | |
134 | // together and configured. This is used to test gRPC connections. | |
135 | func TestPluginGRPCConn(t testing.T, ps map[string]Plugin) (*GRPCClient, *GRPCServer) { | |
136 | // Create a listener | |
137 | l, err := net.Listen("tcp", "127.0.0.1:0") | |
138 | if err != nil { | |
139 | t.Fatalf("err: %s", err) | |
140 | } | |
141 | ||
142 | // Start up the server | |
143 | server := &GRPCServer{ | |
144 | Plugins: ps, | |
107c1cdb | 145 | DoneCh: make(chan struct{}), |
15c0b25d AP |
146 | Server: DefaultGRPCServer, |
147 | Stdout: new(bytes.Buffer), | |
148 | Stderr: new(bytes.Buffer), | |
107c1cdb | 149 | logger: hclog.Default(), |
15c0b25d AP |
150 | } |
151 | if err := server.Init(); err != nil { | |
152 | t.Fatalf("err: %s", err) | |
153 | } | |
154 | go server.Serve(l) | |
155 | ||
156 | // Connect to the server | |
157 | conn, err := grpc.Dial( | |
158 | l.Addr().String(), | |
159 | grpc.WithBlock(), | |
160 | grpc.WithInsecure()) | |
161 | if err != nil { | |
162 | t.Fatalf("err: %s", err) | |
163 | } | |
164 | ||
165 | brokerGRPCClient := newGRPCBrokerClient(conn) | |
166 | broker := newGRPCBroker(brokerGRPCClient, nil) | |
167 | go broker.Run() | |
168 | go brokerGRPCClient.StartStream() | |
169 | ||
170 | // Create the client | |
171 | client := &GRPCClient{ | |
107c1cdb ND |
172 | Conn: conn, |
173 | Plugins: ps, | |
174 | broker: broker, | |
175 | doneCtx: context.Background(), | |
176 | controller: plugin.NewGRPCControllerClient(conn), | |
15c0b25d AP |
177 | } |
178 | ||
179 | return client, server | |
180 | } |