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

   1 package plugin
   2
   3 import (
   4         "bytes"
   5         "net"
   6         "net/rpc"
   7         "testing"
   8 )
   9
  10 // The testing file contains test helpers that you can use outside of
  11 // this package for making it easier to test plugins themselves.
  12
  13 // TestConn is a helper function for returning a client and server
  14 // net.Conn connected to each other.
  15 func TestConn(t *testing.T) (net.Conn, net.Conn) {
  16         // Listen to any local port. This listener will be closed
  17         // after a single connection is established.
  18         l, err := net.Listen("tcp", "127.0.0.1:0")
  19         if err != nil {
  20                 t.Fatalf("err: %s", err)
  21         }
  22
  23         // Start a goroutine to accept our client connection
  24         var serverConn net.Conn
  25         doneCh := make(chan struct{})
  26         go func() {
  27                 defer close(doneCh)
  28                 defer l.Close()
  29                 var err error
  30                 serverConn, err = l.Accept()
  31                 if err != nil {
  32                         t.Fatalf("err: %s", err)
  33                 }
  34         }()
  35
  36         // Connect to the server
  37         clientConn, err := net.Dial("tcp", l.Addr().String())
  38         if err != nil {
  39                 t.Fatalf("err: %s", err)
  40         }
  41
  42         // Wait for the server side to acknowledge it has connected
  43         <-doneCh
  44
  45         return clientConn, serverConn
  46 }
  47
  48 // TestRPCConn returns a rpc client and server connected to each other.
  49 func TestRPCConn(t *testing.T) (*rpc.Client, *rpc.Server) {
  50         clientConn, serverConn := TestConn(t)
  51
  52         server := rpc.NewServer()
  53         go server.ServeConn(serverConn)
  54
  55         client := rpc.NewClient(clientConn)
  56         return client, server
  57 }
  58
  59 // TestPluginRPCConn returns a plugin RPC client and server that are connected
  60 // together and configured.
  61 func TestPluginRPCConn(t *testing.T, ps map[string]Plugin) (*RPCClient, *RPCServer) {
  62         // Create two net.Conns we can use to shuttle our control connection
  63         clientConn, serverConn := TestConn(t)
  64
  65         // Start up the server
  66         server := &RPCServer{Plugins: ps, Stdout: new(bytes.Buffer), Stderr: new(bytes.Buffer)}
  67         go server.ServeConn(serverConn)
  68
  69         // Connect the client to the server
  70         client, err := NewRPCClient(clientConn, ps)
  71         if err != nil {
  72                 t.Fatalf("err: %s", err)
  73         }
  74
  75         return client, server
  76 }