vendor/github.com/hashicorp/go-plugin/testing.go

   1 package plugin
   2
   3 import (
   4         "bytes"
   5         "context"
   6         "io"
   7         "net"
   8         "net/rpc"
   9
  10         "github.com/mitchellh/go-testing-interface"
  11         hclog "github.com/hashicorp/go-hclog"
  12         "github.com/hashicorp/go-plugin/internal/plugin"
  13         "google.golang.org/grpc"
  14 )
  15
  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
  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.
  33 func TestConn(t testing.T) (net.Conn, net.Conn) {
  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.
  67 func TestRPCConn(t testing.T) (*rpc.Client, *rpc.Server) {
  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.
  79 func TestPluginRPCConn(t testing.T, ps map[string]Plugin, opts *TestOptions) (*RPCClient, *RPCServer) {
  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)}
  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         }
  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 }
 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,
 145                 DoneCh:  make(chan struct{}),
 146                 Server:  DefaultGRPCServer,
 147                 Stdout:  new(bytes.Buffer),
 148                 Stderr:  new(bytes.Buffer),
 149                 logger:  hclog.Default(),
 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{
 172                 Conn:       conn,
 173                 Plugins:    ps,
 174                 broker:     broker,
 175                 doneCtx:    context.Background(),
 176                 controller: plugin.NewGRPCControllerClient(conn),
 177         }
 178
 179         return client, server
 180 }