summaryrefslogtreecommitdiffstats
path: root/vendor/github.com/hashicorp/go-plugin/testing.go
blob: 2f541d968399f129c5233c0f6a087cc287645f9d (plain)
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
package plugin

import (
	"bytes"
	"context"
	"io"
	"net"
	"net/rpc"

	"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,
		Server:  DefaultGRPCServer,
		Stdout:  new(bytes.Buffer),
		Stderr:  new(bytes.Buffer),
	}
	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(),
	}

	return client, server
}