github.com/XiaoMi/Gaea@v1.2.5/backend/balancer_test.go (about)

     1  package backend
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/stretchr/testify/assert"
     7  )
     8  
     9  func checkRoundRobin(rb []int, weights []int, gcd int) bool {
    10  	ret := make(map[int]int)
    11  	for _, node := range rb {
    12  		ret[node]++
    13  	}
    14  
    15  	invalidNode := 0
    16  	for node, weight := range weights {
    17  		if weight == 0 {
    18  			invalidNode++
    19  			continue
    20  		}
    21  		if v, ok := ret[node]; !ok {
    22  			return false
    23  		} else if v != weight/gcd {
    24  			return false
    25  		}
    26  	}
    27  
    28  	if len(ret) != len(weights)-invalidNode {
    29  		return false
    30  	}
    31  	return true
    32  }
    33  
    34  func TestBalancer(t *testing.T) {
    35  	weight := []int{2, 1, 4}
    36  	gcd := gcd(weight)
    37  	assert.Equal(t, 1, gcd)
    38  
    39  	b := newBalancer(weight, 4)
    40  	assert.Equal(t, true, checkRoundRobin(b.roundRobinQ, weight, gcd))
    41  }
    42  
    43  func TestBalancerA(t *testing.T) {
    44  	weight := []int{2, 2, 2, 4}
    45  	gcd := gcd(weight)
    46  	assert.Equal(t, 2, gcd)
    47  
    48  	b := newBalancer(weight, 4)
    49  	assert.Equal(t, true, checkRoundRobin(b.roundRobinQ, weight, gcd))
    50  }
    51  
    52  func TestBalancerB(t *testing.T) {
    53  	weight := []int{1, 2}
    54  	gcd := gcd(weight)
    55  	assert.Equal(t, 1, gcd)
    56  
    57  	b := newBalancer(weight, 3)
    58  	assert.Equal(t, true, checkRoundRobin(b.roundRobinQ, weight, gcd))
    59  	assert.Equal(t, 0, b.lastIndex)
    60  
    61  	for i := 0; i < 10; i++ {
    62  		_, err := b.next()
    63  		assert.Equal(t, nil, err)
    64  		assert.Equal(t, (i+1)%len(b.roundRobinQ), b.lastIndex)
    65  	}
    66  }