github.com/diptanu/nomad@v0.5.7-0.20170516172507-d72e86cbe3d9/scheduler/select_test.go (about)

     1  package scheduler
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/hashicorp/nomad/nomad/mock"
     7  )
     8  
     9  func TestLimitIterator(t *testing.T) {
    10  	_, ctx := testContext(t)
    11  	nodes := []*RankedNode{
    12  		&RankedNode{
    13  			Node:  mock.Node(),
    14  			Score: 1,
    15  		},
    16  		&RankedNode{
    17  			Node:  mock.Node(),
    18  			Score: 2,
    19  		},
    20  		&RankedNode{
    21  			Node:  mock.Node(),
    22  			Score: 3,
    23  		},
    24  	}
    25  	static := NewStaticRankIterator(ctx, nodes)
    26  
    27  	limit := NewLimitIterator(ctx, static, 1)
    28  	limit.SetLimit(2)
    29  
    30  	out := collectRanked(limit)
    31  	if len(out) != 2 {
    32  		t.Fatalf("bad: %v", out)
    33  	}
    34  	if out[0] != nodes[0] && out[1] != nodes[1] {
    35  		t.Fatalf("bad: %v", out)
    36  	}
    37  
    38  	out = collectRanked(limit)
    39  	if len(out) != 0 {
    40  		t.Fatalf("bad: %v", out)
    41  	}
    42  	limit.Reset()
    43  
    44  	out = collectRanked(limit)
    45  	if len(out) != 2 {
    46  		t.Fatalf("bad: %v", out)
    47  	}
    48  	if out[0] != nodes[2] && out[1] != nodes[0] {
    49  		t.Fatalf("bad: %v", out)
    50  	}
    51  }
    52  
    53  func TestMaxScoreIterator(t *testing.T) {
    54  	_, ctx := testContext(t)
    55  	nodes := []*RankedNode{
    56  		&RankedNode{
    57  			Node:  mock.Node(),
    58  			Score: 1,
    59  		},
    60  		&RankedNode{
    61  			Node:  mock.Node(),
    62  			Score: 2,
    63  		},
    64  		&RankedNode{
    65  			Node:  mock.Node(),
    66  			Score: 3,
    67  		},
    68  	}
    69  	static := NewStaticRankIterator(ctx, nodes)
    70  
    71  	max := NewMaxScoreIterator(ctx, static)
    72  
    73  	out := collectRanked(max)
    74  	if len(out) != 1 {
    75  		t.Fatalf("bad: %v", out)
    76  	}
    77  	if out[0] != nodes[2] {
    78  		t.Fatalf("bad: %v", out)
    79  	}
    80  
    81  	out = collectRanked(max)
    82  	if len(out) != 0 {
    83  		t.Fatalf("bad: %v", out)
    84  	}
    85  	max.Reset()
    86  
    87  	out = collectRanked(max)
    88  	if len(out) != 1 {
    89  		t.Fatalf("bad: %v", out)
    90  	}
    91  	if out[0] != nodes[2] {
    92  		t.Fatalf("bad: %v", out)
    93  	}
    94  }