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  }