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