github.com/diptanu/nomad@v0.5.7-0.20170516172507-d72e86cbe3d9/nomad/vault_testing.go (about)

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