1 files changed, 82 insertions, 4 deletions

diff --git a/vendor/github.com/hashicorp/go-plugin/testing.go b/vendor/github.com/hashicorp/go-plugin/testing.go
index 9086a1b..df29593 100644
--- a/vendor/github.com/hashicorp/go-plugin/testing.go
+++ b/vendor/github.com/hashicorp/go-plugin/testing.go
@@ -2,9 +2,12 @@ package plugin
 
 import (
         "bytes"
+        "context"
         "net"
         "net/rpc"
-        "testing"
+
+        "github.com/mitchellh/go-testing-interface"
+        "google.golang.org/grpc"
 )
 
 // The testing file contains test helpers that you can use outside of
@@ -12,7 +15,7 @@ import (
 
 // 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) {
+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")
@@ -46,7 +49,7 @@ func TestConn(t *testing.T) (net.Conn, net.Conn) {
 }
 
 // TestRPCConn returns a rpc client and server connected to each other.
-func TestRPCConn(t *testing.T) (*rpc.Client, *rpc.Server) {
+func TestRPCConn(t testing.T) (*rpc.Client, *rpc.Server) {
         clientConn, serverConn := TestConn(t)
 
         server := rpc.NewServer()
@@ -58,7 +61,7 @@ func TestRPCConn(t *testing.T) (*rpc.Client, *rpc.Server) {
 
 // TestPluginRPCConn returns a plugin RPC client and server that are connected
 // together and configured.
-func TestPluginRPCConn(t *testing.T, ps map[string]Plugin) (*RPCClient, *RPCServer) {
+func TestPluginRPCConn(t testing.T, ps map[string]Plugin) (*RPCClient, *RPCServer) {
         // Create two net.Conns we can use to shuttle our control connection
         clientConn, serverConn := TestConn(t)
 
@@ -74,3 +77,78 @@ func TestPluginRPCConn(t *testing.T, ps map[string]Plugin) (*RPCClient, *RPCServ
 
         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
+}