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

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