github.com/kardianos/nomad@v0.1.3-0.20151022182107-b13df73ee850/scheduler/scheduler_test.go (about) 1 package scheduler 2 3 import ( 4 "log" 5 "os" 6 "sync" 7 "testing" 8 9 "github.com/hashicorp/nomad/nomad/state" 10 "github.com/hashicorp/nomad/nomad/structs" 11 ) 12 13 // RejectPlan is used to always reject the entire plan and force a state refresh 14 type RejectPlan struct { 15 Harness *Harness 16 } 17 18 func (r *RejectPlan) SubmitPlan(*structs.Plan) (*structs.PlanResult, State, error) { 19 result := new(structs.PlanResult) 20 result.RefreshIndex = r.Harness.NextIndex() 21 return result, r.Harness.State, nil 22 } 23 24 func (r *RejectPlan) UpdateEval(eval *structs.Evaluation) error { 25 return nil 26 } 27 28 func (r *RejectPlan) CreateEval(*structs.Evaluation) error { 29 return nil 30 } 31 32 // Harness is a lightweight testing harness for schedulers. 33 // It manages a state store copy and provides the planner 34 // interface. It can be extended for various testing uses. 35 type Harness struct { 36 State *state.StateStore 37 38 Planner Planner 39 planLock sync.Mutex 40 41 Plans []*structs.Plan 42 Evals []*structs.Evaluation 43 CreateEvals []*structs.Evaluation 44 45 nextIndex uint64 46 nextIndexLock sync.Mutex 47 } 48 49 // NewHarness is used to make a new testing harness 50 func NewHarness(t *testing.T) *Harness { 51 state, err := state.NewStateStore(os.Stderr) 52 if err != nil { 53 t.Fatalf("err: %v", err) 54 } 55 56 h := &Harness{ 57 State: state, 58 nextIndex: 1, 59 } 60 return h 61 } 62 63 // SubmitPlan is used to handle plan submission 64 func (h *Harness) SubmitPlan(plan *structs.Plan) (*structs.PlanResult, State, error) { 65 // Ensure sequential plan application 66 h.planLock.Lock() 67 defer h.planLock.Unlock() 68 69 // Store the plan 70 h.Plans = append(h.Plans, plan) 71 72 // Check for custom planner 73 if h.Planner != nil { 74 return h.Planner.SubmitPlan(plan) 75 } 76 77 // Get the index 78 index := h.NextIndex() 79 80 // Prepare the result 81 result := new(structs.PlanResult) 82 result.NodeUpdate = plan.NodeUpdate 83 result.NodeAllocation = plan.NodeAllocation 84 result.AllocIndex = index 85 86 // Flatten evicts and allocs 87 var allocs []*structs.Allocation 88 for _, updateList := range plan.NodeUpdate { 89 allocs = append(allocs, updateList...) 90 } 91 for _, allocList := range plan.NodeAllocation { 92 allocs = append(allocs, allocList...) 93 } 94 allocs = append(allocs, plan.FailedAllocs...) 95 96 // Apply the full plan 97 err := h.State.UpsertAllocs(index, allocs) 98 return result, nil, err 99 } 100 101 func (h *Harness) UpdateEval(eval *structs.Evaluation) error { 102 // Ensure sequential plan application 103 h.planLock.Lock() 104 defer h.planLock.Unlock() 105 106 // Store the eval 107 h.Evals = append(h.Evals, eval) 108 109 // Check for custom planner 110 if h.Planner != nil { 111 return h.Planner.UpdateEval(eval) 112 } 113 return nil 114 } 115 116 func (h *Harness) CreateEval(eval *structs.Evaluation) error { 117 // Ensure sequential plan application 118 h.planLock.Lock() 119 defer h.planLock.Unlock() 120 121 // Store the eval 122 h.CreateEvals = append(h.CreateEvals, eval) 123 124 // Check for custom planner 125 if h.Planner != nil { 126 return h.Planner.CreateEval(eval) 127 } 128 return nil 129 } 130 131 // NextIndex returns the next index 132 func (h *Harness) NextIndex() uint64 { 133 h.nextIndexLock.Lock() 134 defer h.nextIndexLock.Unlock() 135 idx := h.nextIndex 136 h.nextIndex += 1 137 return idx 138 } 139 140 // Snapshot is used to snapshot the current state 141 func (h *Harness) Snapshot() State { 142 snap, _ := h.State.Snapshot() 143 return snap 144 } 145 146 // Scheduler is used to return a new scheduler from 147 // a snapshot of current state using the harness for planning. 148 func (h *Harness) Scheduler(factory Factory) Scheduler { 149 logger := log.New(os.Stderr, "", log.LstdFlags) 150 return factory(logger, h.Snapshot(), h) 151 } 152 153 // Process is used to process an evaluation given a factory 154 // function to create the scheduler 155 func (h *Harness) Process(factory Factory, eval *structs.Evaluation) error { 156 sched := h.Scheduler(factory) 157 return sched.Process(eval) 158 } 159 160 func (h *Harness) AssertEvalStatus(t *testing.T, state string) { 161 if len(h.Evals) != 1 { 162 t.Fatalf("bad: %#v", h.Evals) 163 } 164 update := h.Evals[0] 165 166 if update.Status != state { 167 t.Fatalf("bad: %#v", update) 168 } 169 } 170 171 // noErr is used to assert there are no errors 172 func noErr(t *testing.T, err error) { 173 if err != nil { 174 t.Fatalf("err: %v", err) 175 } 176 }