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

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