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 }