github.com/mattyr/nomad@v0.3.3-0.20160919021406-3485a065154a/testutil/vault.go (about)

     1  package testutil
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"os/exec"
     7  	"testing"
     8  
     9  	"github.com/hashicorp/nomad/nomad/structs"
    10  	"github.com/hashicorp/nomad/nomad/structs/config"
    11  	vapi "github.com/hashicorp/vault/api"
    12  )
    13  
    14  // TestVault is a test helper. It uses a fork/exec model to create a test Vault
    15  // server instance in the background and can be initialized with policies, roles
    16  // and backends mounted. The test Vault instances can be used to run a unit test
    17  // and offers and easy API to tear itself down on test end. The only
    18  // prerequisite is that the Vault binary is on the $PATH.
    19  
    20  const (
    21  	// vaultStartPort is the starting port we use to bind Vault servers to
    22  	vaultStartPort uint64 = 40000
    23  )
    24  
    25  // vaultPortOffset is used to atomically increment the port numbers.
    26  var vaultPortOffset uint64
    27  
    28  // TestVault wraps a test Vault server launched in dev mode, suitable for
    29  // testing.
    30  type TestVault struct {
    31  	cmd *exec.Cmd
    32  	t   *testing.T
    33  
    34  	Addr      string
    35  	HTTPAddr  string
    36  	RootToken string
    37  	Config    *config.VaultConfig
    38  	Client    *vapi.Client
    39  }
    40  
    41  // NewTestVault returns a new TestVault instance that has yet to be started
    42  func NewTestVault(t *testing.T) *TestVault {
    43  	port := getPort()
    44  	token := structs.GenerateUUID()
    45  	bind := fmt.Sprintf("-dev-listen-address=127.0.0.1:%d", port)
    46  	http := fmt.Sprintf("http://127.0.0.1:%d", port)
    47  	root := fmt.Sprintf("-dev-root-token-id=%s", token)
    48  
    49  	cmd := exec.Command("vault", "server", "-dev", bind, root)
    50  	cmd.Stdout = os.Stdout
    51  	cmd.Stderr = os.Stderr
    52  
    53  	// Build the config
    54  	conf := vapi.DefaultConfig()
    55  	conf.Address = http
    56  
    57  	// Make the client and set the token to the root token
    58  	client, err := vapi.NewClient(conf)
    59  	if err != nil {
    60  		t.Fatalf("failed to build Vault API client: %v", err)
    61  	}
    62  	client.SetToken(token)
    63  
    64  	tv := &TestVault{
    65  		cmd:       cmd,
    66  		t:         t,
    67  		Addr:      bind,
    68  		HTTPAddr:  http,
    69  		RootToken: token,
    70  		Client:    client,
    71  		Config: &config.VaultConfig{
    72  			Enabled: true,
    73  			Token:   token,
    74  			Addr:    http,
    75  		},
    76  	}
    77  
    78  	return tv
    79  }
    80  
    81  // Start starts the test Vault server and waits for it to respond to its HTTP
    82  // API
    83  func (tv *TestVault) Start() *TestVault {
    84  	if err := tv.cmd.Start(); err != nil {
    85  		tv.t.Fatalf("failed to start vault: %v", err)
    86  	}
    87  
    88  	tv.waitForAPI()
    89  	return tv
    90  }
    91  
    92  // Stop stops the test Vault server
    93  func (tv *TestVault) Stop() {
    94  	if tv.cmd.Process == nil {
    95  		return
    96  	}
    97  
    98  	if err := tv.cmd.Process.Kill(); err != nil {
    99  		tv.t.Errorf("err: %s", err)
   100  	}
   101  	tv.cmd.Wait()
   102  }
   103  
   104  // waitForAPI waits for the Vault HTTP endpoint to start
   105  // responding. This is an indication that the agent has started.
   106  func (tv *TestVault) waitForAPI() {
   107  	WaitForResult(func() (bool, error) {
   108  		inited, err := tv.Client.Sys().InitStatus()
   109  		if err != nil {
   110  			return false, err
   111  		}
   112  		return inited, nil
   113  	}, func(err error) {
   114  		defer tv.Stop()
   115  		tv.t.Fatalf("err: %s", err)
   116  	})
   117  }
   118  
   119  // getPort returns the next available port to bind Vault against
   120  func getPort() uint64 {
   121  	p := vaultStartPort + vaultPortOffset
   122  	vaultPortOffset += 1
   123  	return p
   124  }