-
Notifications
You must be signed in to change notification settings - Fork 1.9k
/
identities_testing.go
82 lines (67 loc) · 2.27 KB
/
identities_testing.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
package consul
import (
"sync"
"github.com/hashicorp/nomad/helper/uuid"
"github.com/hashicorp/nomad/nomad/structs"
)
// MockServiceIdentitiesClient is used for testing the client for managing consul service
// identity tokens.
type MockServiceIdentitiesClient struct {
// deriveTokenErrors maps an allocation ID and tasks to an error when the
// token is derived
deriveTokenErrors map[string]map[string]error
// DeriveTokenFn allows the caller to control the DeriveToken function. If
// not set an error is returned if found in DeriveTokenErrors and otherwise
// a token is generated and returned
DeriveTokenFn TokenDeriverFunc
// lock around everything
lock sync.Mutex
}
var _ ServiceIdentityAPI = (*MockServiceIdentitiesClient)(nil)
// NewMockServiceIdentitiesClient returns a MockServiceIdentitiesClient for testing.
func NewMockServiceIdentitiesClient() *MockServiceIdentitiesClient {
return &MockServiceIdentitiesClient{
deriveTokenErrors: make(map[string]map[string]error),
}
}
func (mtc *MockServiceIdentitiesClient) DeriveSITokens(alloc *structs.Allocation, tasks []string) (map[string]string, error) {
mtc.lock.Lock()
defer mtc.lock.Unlock()
// if the DeriveTokenFn is explicitly set, use that
if mtc.DeriveTokenFn != nil {
return mtc.DeriveTokenFn(alloc, tasks)
}
// generate a token for each task, unless the mock has an error ready for
// one or more of the tasks in which case return that
tokens := make(map[string]string, len(tasks))
for _, task := range tasks {
if m, ok := mtc.deriveTokenErrors[alloc.ID]; ok {
if err, ok := m[task]; ok {
return nil, err
}
}
tokens[task] = uuid.Generate()
}
return tokens, nil
}
func (mtc *MockServiceIdentitiesClient) SetDeriveTokenError(allocID string, tasks []string, err error) {
mtc.lock.Lock()
defer mtc.lock.Unlock()
if _, ok := mtc.deriveTokenErrors[allocID]; !ok {
mtc.deriveTokenErrors[allocID] = make(map[string]error, 10)
}
for _, task := range tasks {
mtc.deriveTokenErrors[allocID][task] = err
}
}
func (mtc *MockServiceIdentitiesClient) DeriveTokenErrors() map[string]map[string]error {
mtc.lock.Lock()
defer mtc.lock.Unlock()
m := make(map[string]map[string]error)
for aID, tasks := range mtc.deriveTokenErrors {
for task, err := range tasks {
m[aID][task] = err
}
}
return m
}