github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/strategy/average_test.go (about)

     1  package strategy
     2  
     3  import (
     4  	"context"
     5  	"sort"
     6  	"testing"
     7  
     8  	"github.com/cockroachdb/errors"
     9  	"github.com/projecteru2/core/types"
    10  	"github.com/stretchr/testify/assert"
    11  )
    12  
    13  func TestAveragePlan(t *testing.T) {
    14  	// 正常的
    15  	nodes := deployedNodes()
    16  	r, err := AveragePlan(context.Background(), nodes, 1, 0, 0)
    17  	assert.NoError(t, err)
    18  	finalCounts := []int{}
    19  	for _, node := range nodes {
    20  		finalCounts = append(finalCounts, node.Count+r[node.Nodename])
    21  	}
    22  	sort.Ints(finalCounts)
    23  	assert.ElementsMatch(t, []int{3, 4, 6, 8}, finalCounts)
    24  
    25  	// nodes len < limit
    26  	nodes = deployedNodes()
    27  	_, err = AveragePlan(context.Background(), nodes, 100, 0, 5)
    28  	assert.Error(t, err)
    29  	// 超过 cap
    30  	nodes = deployedNodes()
    31  	_, err = AveragePlan(context.Background(), nodes, 100, 0, 0)
    32  	assert.Error(t, err)
    33  	assert.True(t, errors.Is(err, types.ErrInsufficientCapacity))
    34  	// 正常 limit
    35  	nodes = deployedNodes()
    36  	_, err = AveragePlan(context.Background(), nodes, 1, 1, 1)
    37  	assert.NoError(t, err)
    38  
    39  	nodes = genNodesByCapCount([]int{1, 2, 3, 4, 5}, []int{3, 3, 3, 3, 3})
    40  	_, err = AveragePlan(context.Background(), nodes, 4, 100, 4)
    41  	assert.Contains(t, err.Error(), "not enough nodes with capacity of 4, require 4 nodes")
    42  
    43  	nodes = genNodesByCapCount([]int{1, 2, 3, 4, 5}, []int{3, 3, 3, 3, 3})
    44  	_, err = AveragePlan(context.Background(), nodes, 2, 100, 0)
    45  	assert.Contains(t, err.Error(), "not enough nodes with capacity of 2, require 5 nodes")
    46  }