github.com/zhizhiboom/nomad@v0.8.5-0.20180907175415-f28fd3a1a056/client/allocrunner/testing.go (about) 1 package allocrunner 2 3 import ( 4 "io/ioutil" 5 "os" 6 "sync" 7 "testing" 8 9 "github.com/boltdb/bolt" 10 "github.com/hashicorp/nomad/client/config" 11 consulApi "github.com/hashicorp/nomad/client/consul" 12 "github.com/hashicorp/nomad/client/vaultclient" 13 "github.com/hashicorp/nomad/helper/testlog" 14 "github.com/hashicorp/nomad/nomad/mock" 15 "github.com/hashicorp/nomad/nomad/structs" 16 ) 17 18 type MockAllocStateUpdater struct { 19 Allocs []*structs.Allocation 20 mu sync.Mutex 21 } 22 23 // Update fulfills the TaskStateUpdater interface 24 func (m *MockAllocStateUpdater) Update(alloc *structs.Allocation) { 25 m.mu.Lock() 26 m.Allocs = append(m.Allocs, alloc) 27 m.mu.Unlock() 28 } 29 30 // Last returns a copy of the last alloc (or nil) sync'd 31 func (m *MockAllocStateUpdater) Last() *structs.Allocation { 32 m.mu.Lock() 33 defer m.mu.Unlock() 34 n := len(m.Allocs) 35 if n == 0 { 36 return nil 37 } 38 return m.Allocs[n-1].Copy() 39 } 40 41 func TestAllocRunnerFromAlloc(t *testing.T, alloc *structs.Allocation, restarts bool) (*MockAllocStateUpdater, *AllocRunner) { 42 conf := config.DefaultConfig() 43 conf.Node = mock.Node() 44 conf.StateDir = os.TempDir() 45 conf.AllocDir = os.TempDir() 46 tmp, _ := ioutil.TempFile("", "state-db") 47 db, _ := bolt.Open(tmp.Name(), 0600, nil) 48 upd := &MockAllocStateUpdater{} 49 if !restarts { 50 *alloc.Job.LookupTaskGroup(alloc.TaskGroup).RestartPolicy = structs.RestartPolicy{Attempts: 0} 51 alloc.Job.Type = structs.JobTypeBatch 52 } 53 vclient := vaultclient.NewMockVaultClient() 54 ar := NewAllocRunner(testlog.Logger(t), conf, db, upd.Update, alloc, vclient, consulApi.NewMockConsulServiceClient(t), NoopPrevAlloc{}) 55 return upd, ar 56 } 57 58 func TestAllocRunner(t *testing.T, restarts bool) (*MockAllocStateUpdater, *AllocRunner) { 59 // Use mock driver 60 alloc := mock.Alloc() 61 task := alloc.Job.TaskGroups[0].Tasks[0] 62 task.Driver = "mock_driver" 63 task.Config["run_for"] = "500ms" 64 return TestAllocRunnerFromAlloc(t, alloc, restarts) 65 }