-
Notifications
You must be signed in to change notification settings - Fork 1.9k
/
testing.go
90 lines (78 loc) · 2.68 KB
/
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
83
84
85
86
87
88
89
90
// +build !release
package allocrunner
import (
"sync"
"testing"
"github.com/hashicorp/nomad/client/lib/cgutil"
"github.com/hashicorp/nomad/client/allocrunner/taskrunner/getter"
"github.com/hashicorp/nomad/client/allocwatcher"
clientconfig "github.com/hashicorp/nomad/client/config"
"github.com/hashicorp/nomad/client/consul"
"github.com/hashicorp/nomad/client/devicemanager"
"github.com/hashicorp/nomad/client/pluginmanager/drivermanager"
"github.com/hashicorp/nomad/client/state"
"github.com/hashicorp/nomad/client/vaultclient"
"github.com/hashicorp/nomad/nomad/structs"
"github.com/stretchr/testify/require"
)
// MockStateUpdater implements the AllocStateHandler interface and records
// alloc updates.
type MockStateUpdater struct {
Updates []*structs.Allocation
mu sync.Mutex
}
// AllocStateUpdated implements the AllocStateHandler interface and records an
// alloc update.
func (m *MockStateUpdater) AllocStateUpdated(alloc *structs.Allocation) {
m.mu.Lock()
m.Updates = append(m.Updates, alloc)
m.mu.Unlock()
}
// Last returns a copy of the last alloc (or nil) update. Safe for concurrent
// access with updates.
func (m *MockStateUpdater) Last() *structs.Allocation {
m.mu.Lock()
defer m.mu.Unlock()
n := len(m.Updates)
if n == 0 {
return nil
}
return m.Updates[n-1].Copy()
}
// Reset resets the recorded alloc updates.
func (m *MockStateUpdater) Reset() {
m.mu.Lock()
m.Updates = nil
m.mu.Unlock()
}
func testAllocRunnerConfig(t *testing.T, alloc *structs.Allocation) (*Config, func()) {
clientConf, cleanup := clientconfig.TestClientConfig(t)
conf := &Config{
// Copy the alloc in case the caller edits and reuses it
Alloc: alloc.Copy(),
Logger: clientConf.Logger,
ClientConfig: clientConf,
StateDB: state.NoopDB{},
Consul: consul.NewMockConsulServiceClient(t, clientConf.Logger),
ConsulSI: consul.NewMockServiceIdentitiesClient(),
Vault: vaultclient.NewMockVaultClient(),
StateUpdater: &MockStateUpdater{},
PrevAllocWatcher: allocwatcher.NoopPrevAlloc{},
PrevAllocMigrator: allocwatcher.NoopPrevAlloc{},
DeviceManager: devicemanager.NoopMockManager(),
DriverManager: drivermanager.TestDriverManager(t),
CpusetManager: cgutil.NoopCpusetManager(),
ServersContactedCh: make(chan struct{}),
Getter: getter.TestDefaultGetter(t),
}
return conf, cleanup
}
func TestAllocRunnerFromAlloc(t *testing.T, alloc *structs.Allocation) (*allocRunner, func()) {
t.Helper()
cfg, cleanup := testAllocRunnerConfig(t, alloc)
ar, err := NewAllocRunner(cfg)
if err != nil {
require.NoError(t, err, "Failed to setup AllocRunner")
}
return ar, cleanup
}