github.com/maier/nomad@v0.4.1-0.20161110003312-a9e3d0b8549d/nomad/vault_testing.go (about)

     1  package nomad
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/hashicorp/nomad/nomad/structs"
     7  	"github.com/hashicorp/nomad/nomad/structs/config"
     8  	vapi "github.com/hashicorp/vault/api"
     9  )
    10  
    11  // TestVaultClient is a Vault client appropriate for use during testing. Its
    12  // behavior is programmable such that endpoints can be tested under various
    13  // circumstances.
    14  type TestVaultClient struct {
    15  	// LookupTokenErrors maps a token to an error that will be returned by the
    16  	// LookupToken call
    17  	LookupTokenErrors map[string]error
    18  
    19  	// LookupTokenSecret maps a token to the Vault secret that will be returned
    20  	// by the LookupToken call
    21  	LookupTokenSecret map[string]*vapi.Secret
    22  
    23  	// CreateTokenErrors maps a token to an error that will be returned by the
    24  	// CreateToken call
    25  	CreateTokenErrors map[string]map[string]error
    26  
    27  	// CreateTokenSecret maps a token to the Vault secret that will be returned
    28  	// by the CreateToken call
    29  	CreateTokenSecret map[string]map[string]*vapi.Secret
    30  
    31  	RevokedTokens []*structs.VaultAccessor
    32  }
    33  
    34  func (v *TestVaultClient) LookupToken(ctx context.Context, token string) (*vapi.Secret, error) {
    35  	var secret *vapi.Secret
    36  	var err error
    37  
    38  	if v.LookupTokenSecret != nil {
    39  		secret = v.LookupTokenSecret[token]
    40  	}
    41  	if v.LookupTokenErrors != nil {
    42  		err = v.LookupTokenErrors[token]
    43  	}
    44  
    45  	return secret, err
    46  }
    47  
    48  // SetLookupTokenSecret sets the error that will be returned by the token
    49  // lookup
    50  func (v *TestVaultClient) SetLookupTokenError(token string, err error) {
    51  	if v.LookupTokenErrors == nil {
    52  		v.LookupTokenErrors = make(map[string]error)
    53  	}
    54  
    55  	v.LookupTokenErrors[token] = err
    56  }
    57  
    58  // SetLookupTokenSecret sets the secret that will be returned by the token
    59  // lookup
    60  func (v *TestVaultClient) SetLookupTokenSecret(token string, secret *vapi.Secret) {
    61  	if v.LookupTokenSecret == nil {
    62  		v.LookupTokenSecret = make(map[string]*vapi.Secret)
    63  	}
    64  
    65  	v.LookupTokenSecret[token] = secret
    66  }
    67  
    68  // SetLookupTokenAllowedPolicies is a helper that adds a secret that allows the
    69  // given policies
    70  func (v *TestVaultClient) SetLookupTokenAllowedPolicies(token string, policies []string) {
    71  	s := &vapi.Secret{
    72  		Data: map[string]interface{}{
    73  			"policies": policies,
    74  		},
    75  	}
    76  
    77  	v.SetLookupTokenSecret(token, s)
    78  }
    79  
    80  func (v *TestVaultClient) CreateToken(ctx context.Context, a *structs.Allocation, task string) (*vapi.Secret, error) {
    81  	var secret *vapi.Secret
    82  	var err error
    83  
    84  	if v.CreateTokenSecret != nil {
    85  		tasks := v.CreateTokenSecret[a.ID]
    86  		if tasks != nil {
    87  			secret = tasks[task]
    88  		}
    89  	}
    90  	if v.CreateTokenErrors != nil {
    91  		tasks := v.CreateTokenErrors[a.ID]
    92  		if tasks != nil {
    93  			err = tasks[task]
    94  		}
    95  	}
    96  
    97  	return secret, err
    98  }
    99  
   100  // SetCreateTokenError sets the error that will be returned by the token
   101  // creation
   102  func (v *TestVaultClient) SetCreateTokenError(allocID, task string, err error) {
   103  	if v.CreateTokenErrors == nil {
   104  		v.CreateTokenErrors = make(map[string]map[string]error)
   105  	}
   106  
   107  	tasks := v.CreateTokenErrors[allocID]
   108  	if tasks == nil {
   109  		tasks = make(map[string]error)
   110  		v.CreateTokenErrors[allocID] = tasks
   111  	}
   112  
   113  	v.CreateTokenErrors[allocID][task] = err
   114  }
   115  
   116  // SetCreateTokenSecret sets the secret that will be returned by the token
   117  // creation
   118  func (v *TestVaultClient) SetCreateTokenSecret(allocID, task string, secret *vapi.Secret) {
   119  	if v.CreateTokenSecret == nil {
   120  		v.CreateTokenSecret = make(map[string]map[string]*vapi.Secret)
   121  	}
   122  
   123  	tasks := v.CreateTokenSecret[allocID]
   124  	if tasks == nil {
   125  		tasks = make(map[string]*vapi.Secret)
   126  		v.CreateTokenSecret[allocID] = tasks
   127  	}
   128  
   129  	v.CreateTokenSecret[allocID][task] = secret
   130  }
   131  
   132  func (v *TestVaultClient) RevokeTokens(ctx context.Context, accessors []*structs.VaultAccessor, committed bool) error {
   133  	v.RevokedTokens = append(v.RevokedTokens, accessors...)
   134  	return nil
   135  }
   136  
   137  func (v *TestVaultClient) Stop()                                      {}
   138  func (v *TestVaultClient) SetActive(enabled bool)                     {}
   139  func (v *TestVaultClient) SetConfig(config *config.VaultConfig) error { return nil }