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  }