github.com/maier/nomad@v0.4.1-0.20161110003312-a9e3d0b8549d/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  	enable := true
    65  	tv := &TestVault{
    66  		cmd:       cmd,
    67  		t:         t,
    68  		Addr:      bind,
    69  		HTTPAddr:  http,
    70  		RootToken: token,
    71  		Client:    client,
    72  		Config: &config.VaultConfig{
    73  			Enabled: &enable,
    74  			Token:   token,
    75  			Addr:    http,
    76  		},
    77  	}
    78  
    79  	return tv
    80  }
    81  
    82  // Start starts the test Vault server and waits for it to respond to its HTTP
    83  // API
    84  func (tv *TestVault) Start() *TestVault {
    85  	if err := tv.cmd.Start(); err != nil {
    86  		tv.t.Fatalf("failed to start vault: %v", err)
    87  	}
    88  
    89  	tv.waitForAPI()
    90  	return tv
    91  }
    92  
    93  // Stop stops the test Vault server
    94  func (tv *TestVault) Stop() {
    95  	if tv.cmd.Process == nil {
    96  		return
    97  	}
    98  
    99  	if err := tv.cmd.Process.Kill(); err != nil {
   100  		tv.t.Errorf("err: %s", err)
   101  	}
   102  	tv.cmd.Wait()
   103  }
   104  
   105  // waitForAPI waits for the Vault HTTP endpoint to start
   106  // responding. This is an indication that the agent has started.
   107  func (tv *TestVault) waitForAPI() {
   108  	WaitForResult(func() (bool, error) {
   109  		inited, err := tv.Client.Sys().InitStatus()
   110  		if err != nil {
   111  			return false, err
   112  		}
   113  		return inited, nil
   114  	}, func(err error) {
   115  		defer tv.Stop()
   116  		tv.t.Fatalf("err: %s", err)
   117  	})
   118  }
   119  
   120  // getPort returns the next available port to bind Vault against
   121  func getPort() uint64 {
   122  	p := vaultStartPort + vaultPortOffset
   123  	vaultPortOffset += 1
   124  	return p
   125  }