github.com/fufuok/balancer@v1.0.0/wr_test.go (about) 1 package balancer 2 3 import ( 4 "sync" 5 "sync/atomic" 6 "testing" 7 ) 8 9 func TestWeightedRand(t *testing.T) { 10 lb := NewWeightedRand() 11 item := lb.Select() 12 if item != "" { 13 t.Fatalf("wr expected empty, actual %s", item) 14 } 15 16 lb.Add("A", 0) 17 item = lb.Select() 18 if item != "" { 19 t.Fatalf("wr expected empty, actual %s", item) 20 } 21 lb.Add("B", 1) 22 item = lb.Select() 23 if item != "B" { 24 t.Fatalf("wr expected B, actual %s", item) 25 } 26 item = lb.Select("test") 27 if item != "B" { 28 t.Fatalf("wr 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 = NewWeightedRand(nodes) 38 count := make(map[string]int) 39 for i := 0; i < 2000; i++ { 40 item := lb.Select() 41 count[item]++ 42 } 43 if count["A"] != 0 || count["B"] <= 150 || count["C"] <= 750 || count["D"] <= 250 { 44 t.Fatal("wr wrong") 45 } 46 if count["A"]+count["B"]+count["C"]+count["D"] != 2000 { 47 t.Fatal("wr wrong") 48 } 49 50 lb.RemoveAll() 51 lb.Add("F", 2) 52 lb.Add("F", 1) 53 all, ok := lb.All().(map[string]int) 54 if !ok || all["F"] != 1 { 55 t.Fatal("wr all() wrong") 56 } 57 58 lb.Remove("F") 59 item = lb.Select() 60 if item != "" { 61 t.Fatalf("wr expected empty, actual %s", item) 62 } 63 64 nodes = map[string]int{ 65 "X": 0, 66 "Y": 1, 67 } 68 ok = lb.Update(nodes) 69 if ok != true { 70 t.Fatal("wr update wrong") 71 } 72 all, ok = lb.All().(map[string]int) 73 if !ok || all["Y"] != 1 { 74 t.Fatal("wr all() wrong") 75 } 76 item = lb.Select() 77 if item != "Y" { 78 t.Fatal("wr update wrong") 79 } 80 } 81 82 func TestWeightedRand_C(t *testing.T) { 83 var ( 84 a, b, c, d int64 85 ) 86 nodes := map[string]int{ 87 "A": 5, 88 "B": 1, 89 "C": 4, 90 "D": 0, 91 } 92 lb := NewWeightedRand(nodes) 93 94 var wg sync.WaitGroup 95 for i := 0; i < 500; i++ { 96 wg.Add(1) 97 go func() { 98 defer wg.Done() 99 for j := 0; j < 2000; j++ { 100 switch lb.Select() { 101 case "A": 102 atomic.AddInt64(&a, 1) 103 case "B": 104 atomic.AddInt64(&b, 1) 105 case "C": 106 atomic.AddInt64(&c, 1) 107 case "D": 108 atomic.AddInt64(&d, 1) 109 } 110 } 111 }() 112 } 113 wg.Wait() 114 115 if atomic.LoadInt64(&a) <= 410000 { 116 t.Fatal("wr wrong: a") 117 } 118 if atomic.LoadInt64(&b) <= 90000 { 119 t.Fatal("wr wrong: b") 120 } 121 if atomic.LoadInt64(&c) <= 390000 { 122 t.Fatal("wr wrong: c") 123 } 124 if atomic.LoadInt64(&d) != 0 { 125 t.Fatal("wr wrong: d") 126 } 127 if atomic.LoadInt64(&a)+atomic.LoadInt64(&b)+atomic.LoadInt64(&c) != 1000000 { 128 t.Fatal("wr wrong: sum") 129 } 130 }