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