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  }