shuffle_test.go

package sharding

import (
	"os"
	"testing"

	"github.com/argoproj/argo-cd/v2/common"
	"github.com/argoproj/argo-cd/v2/pkg/apis/application/v1alpha1"
	dbmocks "github.com/argoproj/argo-cd/v2/util/db/mocks"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
)

func TestShuffle(t *testing.T) {
	db := dbmocks.ArgoDB{}
	cluster1 := createCluster("cluster1", "10")
	cluster2 := createCluster("cluster2", "20")
	cluster3 := createCluster("cluster3", "30")
	cluster4 := createCluster("cluster4", "40")
	cluster5 := createCluster("cluster5", "50")
	cluster6 := createCluster("cluster6", "60")
	cluster25 := createCluster("cluster6", "25")

	clusterList := &v1alpha1.ClusterList{Items: []v1alpha1.Cluster{cluster1, cluster2, cluster3, cluster4, cluster5, cluster6}}
	db.On("ListClusters", mock.Anything).Return(clusterList, nil)

	// Test with replicas set to 3
	os.Setenv(common.EnvControllerReplicas, "3")
	distributionFunction := GetShardByIndexModuloReplicasCountDistributionFunction(&db)
	assert.Equal(t, 0, distributionFunction(nil))
	assert.Equal(t, 0, distributionFunction(&cluster1))
	assert.Equal(t, 1, distributionFunction(&cluster2))
	assert.Equal(t, 2, distributionFunction(&cluster3))
	assert.Equal(t, 0, distributionFunction(&cluster4))
	assert.Equal(t, 1, distributionFunction(&cluster5))
	assert.Equal(t, 2, distributionFunction(&cluster6))

	// Now, we remove cluster1, it should be unassigned, and all the other should be resuffled
	clusterList.Items = remove(clusterList.Items, 0)
	assert.Equal(t, -1, distributionFunction(&cluster1))
	assert.Equal(t, 0, distributionFunction(&cluster2))
	assert.Equal(t, 1, distributionFunction(&cluster3))
	assert.Equal(t, 2, distributionFunction(&cluster4))
	assert.Equal(t, 0, distributionFunction(&cluster5))
	assert.Equal(t, 1, distributionFunction(&cluster6))

	// Now, we add a cluster with an id=25 so it will be placed right after cluster2
	clusterList.Items = append(clusterList.Items, cluster25)
	assert.Equal(t, -1, distributionFunction(&cluster1))
	assert.Equal(t, 0, distributionFunction(&cluster2))
	assert.Equal(t, 1, distributionFunction(&cluster25))
	assert.Equal(t, 2, distributionFunction(&cluster3))
	assert.Equal(t, 0, distributionFunction(&cluster4))
	assert.Equal(t, 1, distributionFunction(&cluster5))
	assert.Equal(t, 2, distributionFunction(&cluster6))

}

func remove(slice []v1alpha1.Cluster, s int) []v1alpha1.Cluster {
	return append(slice[:s], slice[s+1:]...)
}