github.com/fufuok/balancer@v1.0.0/wrr_test.go (about)

     1  package balancer
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  	"testing"
     7  )
     8  
     9  func TestWeightedRoundRobin(t *testing.T) {
    10  	lb := NewWeightedRoundRobin()
    11  	item := lb.Select()
    12  	if item != "" {
    13  		t.Fatalf("wrr expected empty, actual %s", item)
    14  	}
    15  
    16  	lb.Add("A", 0)
    17  	item = lb.Select()
    18  	if item != "" {
    19  		t.Fatalf("wrr expected empty, actual %s", item)
    20  	}
    21  	lb.Add("B", 1)
    22  	item = lb.Select()
    23  	if item != "B" {
    24  		t.Fatalf("wrr expected B, actual %s", item)
    25  	}
    26  	item = lb.Select("test")
    27  	if item != "B" {
    28  		t.Fatalf("wrr expected B, actual %s", item)
    29  	}
    30  
    31  	nodes := map[string]int{
    32  		"A": 0,
    33  		"B": 1,
    34  		"C": 7,
    35  		"D": 2,
    36  	}
    37  	lb = NewWeightedRoundRobin(nodes)
    38  	count := make(map[string]int)
    39  	for i := 0; i < 1000; i++ {
    40  		item := lb.Select()
    41  		count[item]++
    42  	}
    43  	if count["A"] != 0 || count["B"] != 100 || count["C"] != 700 || count["D"] != 200 {
    44  		t.Fatal("wrr wrong")
    45  	}
    46  
    47  	lb.Reset()
    48  
    49  	lb.Add("E", 10)
    50  	for i := 0; i < 2000; i++ {
    51  		item := lb.Select()
    52  		count[item]++
    53  	}
    54  	if count["A"] != 0 || count["B"] != 200 || count["C"] != 1400 || count["D"] != 400 || count["E"] != 1000 {
    55  		t.Fatal("wrr reset() wrong")
    56  	}
    57  
    58  	ok := lb.Remove("E")
    59  	if ok != true {
    60  		t.Fatal("wrr remove() wrong")
    61  	}
    62  
    63  	lb.Reset()
    64  
    65  	for i := 0; i < 1000; i++ {
    66  		item := lb.Select()
    67  		count[item]++
    68  	}
    69  	if count["A"] != 0 || count["B"] != 300 || count["C"] != 2100 || count["D"] != 600 {
    70  		t.Fatal("wrr wrong")
    71  	}
    72  
    73  	lb.RemoveAll()
    74  	lb.Add("F", 2)
    75  	lb.Add("F", 1)
    76  	all, ok := lb.All().(map[string]int)
    77  	if !ok || all["F"] != 1 {
    78  		t.Fatal("wrr all() wrong")
    79  	}
    80  
    81  	nodes = map[string]int{
    82  		"X": 0,
    83  		"Y": 1,
    84  	}
    85  	ok = lb.Update(nodes)
    86  	if ok != true {
    87  		t.Fatal("wrr update wrong")
    88  	}
    89  	all, ok = lb.All().(map[string]int)
    90  	if !ok || all["Y"] != 1 {
    91  		t.Fatal("wrr all() wrong")
    92  	}
    93  	item = lb.Select()
    94  	if item != "Y" {
    95  		t.Fatal("wrr update wrong")
    96  	}
    97  }
    98  
    99  func TestWeightedRoundRobin_C(t *testing.T) {
   100  	var (
   101  		a, b, c, d int64
   102  	)
   103  	nodes := map[string]int{
   104  		"A": 5,
   105  		"B": 1,
   106  		"C": 4,
   107  		"D": 0,
   108  	}
   109  	lb := NewWeightedRoundRobin(nodes)
   110  
   111  	var wg sync.WaitGroup
   112  	for i := 0; i < 500; i++ {
   113  		wg.Add(1)
   114  		go func() {
   115  			defer wg.Done()
   116  			for j := 0; j < 2000; j++ {
   117  				switch lb.Select() {
   118  				case "A":
   119  					atomic.AddInt64(&a, 1)
   120  				case "B":
   121  					atomic.AddInt64(&b, 1)
   122  				case "C":
   123  					atomic.AddInt64(&c, 1)
   124  				case "D":
   125  					atomic.AddInt64(&d, 1)
   126  				}
   127  			}
   128  		}()
   129  	}
   130  	wg.Wait()
   131  
   132  	if atomic.LoadInt64(&a) != 500000 {
   133  		t.Fatal("wrr wrong: a")
   134  	}
   135  	if atomic.LoadInt64(&b) != 100000 {
   136  		t.Fatal("wrr wrong: b")
   137  	}
   138  	if atomic.LoadInt64(&c) != 400000 {
   139  		t.Fatal("wrr wrong: c")
   140  	}
   141  	if atomic.LoadInt64(&d) != 0 {
   142  		t.Fatal("wrr wrong: d")
   143  	}
   144  }