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

     1  package balancer
     2  
     3  import (
     4  	"strings"
     5  	"sync"
     6  	"sync/atomic"
     7  	"testing"
     8  )
     9  
    10  func TestRandom(t *testing.T) {
    11  	lb := NewRandom()
    12  	item := lb.Select()
    13  	if item != "" {
    14  		t.Fatalf("r expected empty, actual %s", item)
    15  	}
    16  
    17  	lb.Add("A")
    18  	item = lb.Select()
    19  	if item != "A" {
    20  		t.Fatalf("r expected A, actual %s", item)
    21  	}
    22  
    23  	nodes := []string{"A", "B", "C", "D"}
    24  	lb = NewRandom(nodes)
    25  	count := make(map[string]int)
    26  	for i := 0; i < 2000; i++ {
    27  		item := lb.Select()
    28  		count[item]++
    29  	}
    30  	if count["A"] <= 300 || count["B"] <= 300 || count["C"] <= 300 || count["D"] <= 300 {
    31  		t.Fatal("r wrong")
    32  	}
    33  	if count["A"]+count["B"]+count["C"]+count["D"] != 2000 {
    34  		t.Fatal("r wrong")
    35  	}
    36  
    37  	lb.Add("A")
    38  	lb.Add("A")
    39  	lb.Add("C")
    40  	lb.Add("A")
    41  	all := lb.All().([]string)
    42  	if strings.Join(all, "") != "ABCDAACA" {
    43  		t.Fatal("r all() wrong")
    44  	}
    45  
    46  	lb.Remove("C")
    47  	all = lb.All().([]string)
    48  	if strings.Join(all, "") != "ABDAACA" {
    49  		t.Fatal("r all() wrong")
    50  	}
    51  
    52  	lb.Remove("A", true)
    53  	all = lb.All().([]string)
    54  	if strings.Join(all, "") != "BDC" {
    55  		t.Fatal("r all() wrong")
    56  	}
    57  
    58  	lb.RemoveAll()
    59  	lb.Add("F", 1)
    60  	all, ok := lb.All().([]string)
    61  	if !ok || len(all) != 1 {
    62  		t.Fatal("r all() wrong")
    63  	}
    64  
    65  	nodes = []string{"X", "Y"}
    66  	ok = lb.Update(nodes)
    67  	if ok != true {
    68  		t.Fatal("r update wrong")
    69  	}
    70  	item = lb.Select()
    71  	if item != "X" && item != "Y" {
    72  		t.Fatal("r update wrong")
    73  	}
    74  }
    75  
    76  func TestRandom_C(t *testing.T) {
    77  	var (
    78  		a, b, c, d int64
    79  	)
    80  	nodes := []string{"A", "B", "C", "D"}
    81  	lb := NewRandom(nodes)
    82  
    83  	var wg sync.WaitGroup
    84  	for i := 0; i < 500; i++ {
    85  		wg.Add(1)
    86  		go func() {
    87  			defer wg.Done()
    88  			for j := 0; j < 2000; j++ {
    89  				switch lb.Select() {
    90  				case "A":
    91  					atomic.AddInt64(&a, 1)
    92  				case "B":
    93  					atomic.AddInt64(&b, 1)
    94  				case "C":
    95  					atomic.AddInt64(&c, 1)
    96  				case "D":
    97  					atomic.AddInt64(&d, 1)
    98  				}
    99  			}
   100  		}()
   101  	}
   102  	wg.Wait()
   103  
   104  	if atomic.LoadInt64(&a) <= 200000 {
   105  		t.Fatal("r wrong: a")
   106  	}
   107  	if atomic.LoadInt64(&b) <= 200000 {
   108  		t.Fatal("r wrong: b")
   109  	}
   110  	if atomic.LoadInt64(&c) <= 200000 {
   111  		t.Fatal("r wrong: c")
   112  	}
   113  	if atomic.LoadInt64(&d) <= 200000 {
   114  		t.Fatal("r wrong: d")
   115  	}
   116  	if atomic.LoadInt64(&a)+atomic.LoadInt64(&b)+atomic.LoadInt64(&c)+atomic.LoadInt64(&d) != 1000000 {
   117  		t.Fatal("r wrong: sum")
   118  	}
   119  }