github.com/hhrutter/nomad@v0.6.0-rc2.0.20170723054333-80c4b03f0705/scheduler/testing.go (about) 1 package scheduler 2 3 import ( 4 "fmt" 5 "log" 6 "os" 7 "sync" 8 "testing" 9 "time" 10 11 memdb "github.com/hashicorp/go-memdb" 12 "github.com/hashicorp/nomad/nomad/state" 13 "github.com/hashicorp/nomad/nomad/structs" 14 ) 15 16 // RejectPlan is used to always reject the entire plan and force a state refresh 17 type RejectPlan struct { 18 Harness *Harness 19 } 20 21 func (r *RejectPlan) SubmitPlan(*structs.Plan) (*structs.PlanResult, State, error) { 22 result := new(structs.PlanResult) 23 result.RefreshIndex = r.Harness.NextIndex() 24 return result, r.Harness.State, nil 25 } 26 27 func (r *RejectPlan) UpdateEval(eval *structs.Evaluation) error { 28 return nil 29 } 30 31 func (r *RejectPlan) CreateEval(*structs.Evaluation) error { 32 return nil 33 } 34 35 func (r *RejectPlan) ReblockEval(*structs.Evaluation) error { 36 return nil 37 } 38 39 // Harness is a lightweight testing harness for schedulers. It manages a state 40 // store copy and provides the planner interface. It can be extended for various 41 // testing uses or for invoking the scheduler without side effects. 42 type Harness struct { 43 State *state.StateStore 44 45 Planner Planner 46 planLock sync.Mutex 47 48 Plans []*structs.Plan 49 Evals []*structs.Evaluation 50 CreateEvals []*structs.Evaluation 51 ReblockEvals []*structs.Evaluation 52 53 nextIndex uint64 54 nextIndexLock sync.Mutex 55 } 56 57 // NewHarness is used to make a new testing harness 58 func NewHarness(t *testing.T) *Harness { 59 state, err := state.NewStateStore(os.Stderr) 60 if err != nil { 61 t.Fatalf("err: %v", err) 62 } 63 64 h := &Harness{ 65 State: state, 66 nextIndex: 1, 67 } 68 return h 69 } 70 71 // NewHarnessWithState creates a new harness with the given state for testing 72 // purposes. 73 func NewHarnessWithState(t *testing.T, state *state.StateStore) *Harness { 74 return &Harness{ 75 State: state, 76 nextIndex: 1, 77 } 78 } 79 80 // SubmitPlan is used to handle plan submission 81 func (h *Harness) SubmitPlan(plan *structs.Plan) (*structs.PlanResult, State, error) { 82 // Ensure sequential plan application 83 h.planLock.Lock() 84 defer h.planLock.Unlock() 85 86 // Store the plan 87 h.Plans = append(h.Plans, plan) 88 89 // Check for custom planner 90 if h.Planner != nil { 91 return h.Planner.SubmitPlan(plan) 92 } 93 94 // Get the index 95 index := h.NextIndex() 96 97 // Prepare the result 98 result := new(structs.PlanResult) 99 result.NodeUpdate = plan.NodeUpdate 100 result.NodeAllocation = plan.NodeAllocation 101 result.AllocIndex = index 102 103 // Flatten evicts and allocs 104 var allocs []*structs.Allocation 105 for _, updateList := range plan.NodeUpdate { 106 allocs = append(allocs, updateList...) 107 } 108 for _, allocList := range plan.NodeAllocation { 109 allocs = append(allocs, allocList...) 110 } 111 112 // Set the time the alloc was applied for the first time. This can be used 113 // to approximate the scheduling time. 114 now := time.Now().UTC().UnixNano() 115 for _, alloc := range allocs { 116 if alloc.CreateTime == 0 { 117 alloc.CreateTime = now 118 } 119 } 120 121 // Setup the update request 122 req := structs.ApplyPlanResultsRequest{ 123 AllocUpdateRequest: structs.AllocUpdateRequest{ 124 Job: plan.Job, 125 Alloc: allocs, 126 }, 127 Deployment: plan.Deployment, 128 DeploymentUpdates: plan.DeploymentUpdates, 129 } 130 131 // Apply the full plan 132 err := h.State.UpsertPlanResults(index, &req) 133 return result, nil, err 134 } 135 136 func (h *Harness) UpdateEval(eval *structs.Evaluation) error { 137 // Ensure sequential plan application 138 h.planLock.Lock() 139 defer h.planLock.Unlock() 140 141 // Store the eval 142 h.Evals = append(h.Evals, eval) 143 144 // Check for custom planner 145 if h.Planner != nil { 146 return h.Planner.UpdateEval(eval) 147 } 148 return nil 149 } 150 151 func (h *Harness) CreateEval(eval *structs.Evaluation) error { 152 // Ensure sequential plan application 153 h.planLock.Lock() 154 defer h.planLock.Unlock() 155 156 // Store the eval 157 h.CreateEvals = append(h.CreateEvals, eval) 158 159 // Check for custom planner 160 if h.Planner != nil { 161 return h.Planner.CreateEval(eval) 162 } 163 return nil 164 } 165 166 func (h *Harness) ReblockEval(eval *structs.Evaluation) error { 167 // Ensure sequential plan application 168 h.planLock.Lock() 169 defer h.planLock.Unlock() 170 171 // Check that the evaluation was already blocked. 172 ws := memdb.NewWatchSet() 173 old, err := h.State.EvalByID(ws, eval.ID) 174 if err != nil { 175 return err 176 } 177 178 if old == nil { 179 return fmt.Errorf("evaluation does not exist to be reblocked") 180 } 181 if old.Status != structs.EvalStatusBlocked { 182 return fmt.Errorf("evaluation %q is not already in a blocked state", old.ID) 183 } 184 185 h.ReblockEvals = append(h.ReblockEvals, eval) 186 return nil 187 } 188 189 // NextIndex returns the next index 190 func (h *Harness) NextIndex() uint64 { 191 h.nextIndexLock.Lock() 192 defer h.nextIndexLock.Unlock() 193 idx := h.nextIndex 194 h.nextIndex += 1 195 return idx 196 } 197 198 // Snapshot is used to snapshot the current state 199 func (h *Harness) Snapshot() State { 200 snap, _ := h.State.Snapshot() 201 return snap 202 } 203 204 // Scheduler is used to return a new scheduler from 205 // a snapshot of current state using the harness for planning. 206 func (h *Harness) Scheduler(factory Factory) Scheduler { 207 logger := log.New(os.Stderr, "", log.LstdFlags) 208 return factory(logger, h.Snapshot(), h) 209 } 210 211 // Process is used to process an evaluation given a factory 212 // function to create the scheduler 213 func (h *Harness) Process(factory Factory, eval *structs.Evaluation) error { 214 sched := h.Scheduler(factory) 215 return sched.Process(eval) 216 } 217 218 func (h *Harness) AssertEvalStatus(t *testing.T, state string) { 219 if len(h.Evals) != 1 { 220 t.Fatalf("bad: %#v", h.Evals) 221 } 222 update := h.Evals[0] 223 224 if update.Status != state { 225 t.Fatalf("bad: %#v", update) 226 } 227 }