[flytepropeller] Add test under flytepropeller/manager (#4237)

Signed-off-by: Chao-Heng Lee <[email protected]>

flytepropeller/manager/config/config_test.go

@@ -0,0 +1,11 @@
+package config
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetConfig(t *testing.T) {
+	assert.Equal(t, DefaultConfig, GetConfig())
+}

flytepropeller/manager/manager_test.go

@@ -3,17 +3,24 @@ package manager
 import (
 	"context"
 	"fmt"
+	"reflect"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/client-go/kubernetes/fake"
+	"k8s.io/client-go/tools/leaderelection"
 
+	managerConfig "github.com/flyteorg/flyte/flytepropeller/manager/config"
 	"github.com/flyteorg/flyte/flytepropeller/manager/shardstrategy"
 	"github.com/flyteorg/flyte/flytepropeller/manager/shardstrategy/mocks"
+	propellerConfig "github.com/flyteorg/flyte/flytepropeller/pkg/controller/config"
+	leader "github.com/flyteorg/flyte/flytepropeller/pkg/leaderelection"
+	"github.com/flyteorg/flyte/flytestdlib/config"
 	"github.com/flyteorg/flyte/flytestdlib/promutils"
 )
 
@@ -195,3 +202,148 @@ func TestGetPodNames(t *testing.T) {
 		})
 	}
 }
+
+func createPropellerConfig(enableLeaderElection bool) *propellerConfig.Config {
+	return &propellerConfig.Config{
+		LeaderElection: propellerConfig.LeaderElectionConfig{
+			Enabled:       enableLeaderElection,
+			LeaseDuration: config.Duration{Duration: time.Second * 15},
+			RenewDeadline: config.Duration{Duration: time.Second * 10},
+			RetryPeriod:   config.Duration{Duration: time.Second * 2},
+		},
+	}
+}
+
+func TestRun(t *testing.T) {
+	t.Parallel()
+	// setup for leaderElector
+	setup := func() *leaderelection.LeaderElector {
+		kubeClient := fake.NewSimpleClientset(podTemplate)
+		propellerCfg := createPropellerConfig(true)
+		lock, _ := leader.NewResourceLock(kubeClient.CoreV1(), kubeClient.CoordinationV1(), nil, propellerCfg.LeaderElection)
+		leConfig := leaderelection.LeaderElectionConfig{
+			Lock:          lock,
+			LeaseDuration: time.Second * 15,
+			RenewDeadline: time.Second * 10,
+			RetryPeriod:   time.Second * 2,
+			Callbacks: leaderelection.LeaderCallbacks{
+				OnStartedLeading: func(ctx context.Context) {},
+				OnStoppedLeading: func() {},
+				OnNewLeader:      func(identity string) {},
+			},
+		}
+		leaderElector, _ := leaderelection.NewLeaderElector(leConfig)
+		return leaderElector
+	}
+
+	tests := []struct {
+		name          string
+		leaderElector *leaderelection.LeaderElector
+	}{
+		{"withLeaderElection", setup()},
+		{"withoutLeaderElection", nil},
+	}
+
+	metrics := newManagerMetrics(promutils.NewScope(fmt.Sprintf("create_%s", "Run")))
+	kubeClient := fake.NewSimpleClientset(podTemplate)
+	shardStrategy := createShardStrategy(2)
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.TODO())
+			defer cancel()
+
+			manager := Manager{
+				leaderElector:  tt.leaderElector,
+				kubeClient:     kubeClient,
+				metrics:        metrics,
+				podApplication: "flytepropeller",
+				shardStrategy:  shardStrategy,
+			}
+
+			errChan := make(chan error, 1)
+			go func() {
+				err := manager.Run(ctx)
+				errChan <- err
+			}()
+
+			time.Sleep(10 * time.Millisecond)
+			cancel()
+
+			err := <-errChan
+			assert.NoError(t, err)
+		})
+	}
+}
+
+func Test_Run(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.TODO())
+	defer cancel()
+	scope := promutils.NewScope(fmt.Sprintf("create_%s", "run"))
+	kubeClient := fake.NewSimpleClientset(podTemplate)
+	shardStrategy := createShardStrategy(2)
+
+	manager := Manager{
+		kubeClient:     kubeClient,
+		metrics:        newManagerMetrics(scope),
+		podApplication: "flytepropeller",
+		shardStrategy:  shardStrategy,
+	}
+
+	errChan := make(chan error, 1)
+
+	go func() {
+		err := manager.run(ctx)
+		errChan <- err
+	}()
+
+	time.Sleep(10 * time.Millisecond)
+	cancel()
+
+	err := <-errChan
+	assert.NoError(t, err)
+}
+
+func TestNewManager(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name                 string
+		propellerConfig      *propellerConfig.Config
+		leaderElectionEnable bool
+	}{
+		{"enableLeaderElection", createPropellerConfig(true), true},
+		{"withoutLeaderElection", createPropellerConfig(false), false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := context.TODO()
+			managerCfg := managerConfig.DefaultConfig
+			podNamespace := "flyte"
+			ownerReference := make([]metav1.OwnerReference, 0)
+			kubeClient := fake.NewSimpleClientset(podTemplate)
+			scope := promutils.NewScope(fmt.Sprintf("create_%s", tt.name))
+			shardStrategy, _ := shardstrategy.NewShardStrategy(ctx, managerCfg.ShardConfig)
+
+			actualManager, err := New(ctx, tt.propellerConfig, managerCfg, podNamespace, ownerReference, kubeClient, scope)
+
+			expectedManager := &Manager{
+				kubeClient:               kubeClient,
+				leaderElector:            actualManager.leaderElector,
+				metrics:                  actualManager.metrics,
+				ownerReferences:          ownerReference,
+				podApplication:           managerCfg.PodApplication,
+				podNamespace:             podNamespace,
+				podTemplateContainerName: managerCfg.PodTemplateContainerName,
+				podTemplateName:          managerCfg.PodTemplateName,
+				podTemplateNamespace:     managerCfg.PodTemplateNamespace,
+				scanInterval:             managerCfg.ScanInterval.Duration,
+				shardStrategy:            shardStrategy,
+			}
+
+			assert.NoError(t, err)
+			assert.True(t, reflect.DeepEqual(expectedManager, actualManager))
+			assert.Equal(t, scope, actualManager.metrics.Scope)
+			assert.Equal(t, tt.leaderElectionEnable, actualManager.leaderElector != nil)
+		})
+	}
+}

flytepropeller/manager/shardstrategy/environment_test.go

@@ -0,0 +1,31 @@
+package shardstrategy
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestEnvironmentHashCode(t *testing.T) {
+	expectedStrategy := &EnvironmentShardStrategy{
+		EnvType: Project,
+		PerShardIDs: [][]string{
+			{"flytesnacks"},
+			{"flytefoo", "flytebar"},
+			{"*"},
+		},
+	}
+	actualStrategy := &EnvironmentShardStrategy{
+		EnvType: Project,
+		PerShardIDs: [][]string{
+			{"flytesnacks"},
+			{"flytefoo", "flytebar"},
+			{"*"},
+		},
+	}
+	expectedHashcode, err := expectedStrategy.HashCode()
+	assert.NoError(t, err)
+	actualHashcode, err := actualStrategy.HashCode()
+	assert.NoError(t, err)
+	assert.Equal(t, expectedHashcode, actualHashcode)
+}

flytepropeller/manager/shardstrategy/hash_test.go

@@ -23,3 +23,17 @@ func TestComputeKeyRange(t *testing.T) {
 		assert.Equal(t, keyspaceSize, keysCovered)
 	}
 }
+
+func TestHashHashCode(t *testing.T) {
+	expectedStrategy := &HashShardStrategy{
+		ShardCount: 3,
+	}
+	actualStrategy := &HashShardStrategy{
+		ShardCount: 3,
+	}
+	expectedHashcode, err := expectedStrategy.HashCode()
+	assert.NoError(t, err)
+	actualHashcode, err := actualStrategy.HashCode()
+	assert.NoError(t, err)
+	assert.Equal(t, expectedHashcode, actualHashcode)
+}

flytepropeller/manager/shardstrategy/shard_strategy_test.go

@@ -1,10 +1,14 @@
 package shardstrategy
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	v1 "k8s.io/api/core/v1"
+
+	"github.com/flyteorg/flyte/flytepropeller/manager/config"
+	"github.com/flyteorg/flyte/flytepropeller/pkg/apis/flyteworkflow/v1alpha1"
 )
 
 var (
@@ -158,3 +162,141 @@ func TestUpdatePodSpecInvalidPodSpec(t *testing.T) {
 		})
 	}
 }
+
+var (
+	hashShardConfig = config.ShardConfig{
+		Type:       config.ShardTypeHash,
+		ShardCount: 3,
+	}
+	projectShardConfig = config.ShardConfig{
+		Type: config.ShardTypeProject,
+		PerShardMappings: []config.PerShardMappingsConfig{
+			{IDs: []string{"flytesnacks"}},
+			{IDs: []string{"flytefoo", "flytebar"}},
+		},
+	}
+	projectShardWildcardConfig = config.ShardConfig{
+		Type: config.ShardTypeProject,
+		PerShardMappings: []config.PerShardMappingsConfig{
+			{IDs: []string{"flytesnacks"}},
+			{IDs: []string{"flytefoo", "flytebar"}},
+			{IDs: []string{"*"}},
+		},
+	}
+	domainShardConfig = config.ShardConfig{
+		Type: config.ShardTypeDomain,
+		PerShardMappings: []config.PerShardMappingsConfig{
+			{IDs: []string{"production"}},
+			{IDs: []string{"foo", "bar"}},
+		},
+	}
+	domainShardWildcardConfig = config.ShardConfig{
+		Type: config.ShardTypeDomain,
+		PerShardMappings: []config.PerShardMappingsConfig{
+			{IDs: []string{"production"}},
+			{IDs: []string{"foo", "bar"}},
+			{IDs: []string{"*"}},
+		},
+	}
+)
+
+func TestNewShardStrategy(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name          string
+		ShardStrategy ShardStrategy
+		shardConfig   config.ShardConfig
+	}{
+		{"hash", hashShardStrategy, hashShardConfig},
+		{"project", projectShardStrategy, projectShardConfig},
+		{"project_wildcard", projectShardStrategyWildcard, projectShardWildcardConfig},
+		{"domain", domainShardStrategy, domainShardConfig},
+		{"domain_wildcard", domainShardStrategyWildcard, domainShardWildcardConfig},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			strategy, err := NewShardStrategy(context.TODO(), tt.shardConfig)
+			assert.NoError(t, err)
+			assert.Equal(t, tt.ShardStrategy, strategy)
+		})
+	}
+}
+
+var (
+	errorHashShardConfig1 = config.ShardConfig{
+		Type:       config.ShardTypeHash,
+		ShardCount: 0,
+	}
+	errorHashShardConfig2 = config.ShardConfig{
+		Type:       config.ShardTypeHash,
+		ShardCount: v1alpha1.ShardKeyspaceSize + 1,
+	}
+	errorProjectShardConfig = config.ShardConfig{
+		Type: config.ShardTypeProject,
+		PerShardMappings: []config.PerShardMappingsConfig{
+			{IDs: []string{}},
+		},
+	}
+	errorProjectShardWildcardConfig = config.ShardConfig{
+		Type: config.ShardTypeProject,
+		PerShardMappings: []config.PerShardMappingsConfig{
+			{IDs: []string{"flytesnacks", "*"}},
+		},
+	}
+	errorDomainShardWildcardConfig = config.ShardConfig{
+		Type: config.ShardTypeDomain,
+		PerShardMappings: []config.PerShardMappingsConfig{
+			{IDs: []string{"production"}},
+			{IDs: []string{"*"}},
+			{IDs: []string{"*"}},
+		},
+	}
+	errorShardType = config.ShardConfig{
+		Type: -1,
+	}
+)
+
+func TestNewShardStrategyErrorConfig(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name        string
+		shardConfig config.ShardConfig
+	}{
+		{"hash_shard_cnt_zero", errorHashShardConfig1},
+		{"hash_shard_cnt_larger_than_keyspace", errorHashShardConfig2},
+		{"project_with_zero_config_ids", errorProjectShardConfig},
+		{"project_wildcard_with_other_ids", errorProjectShardWildcardConfig},
+		{"domain_multi_wildcard_ids", errorDomainShardWildcardConfig},
+		{"error_shard_type", errorShardType},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			_, err := NewShardStrategy(context.TODO(), tt.shardConfig)
+			assert.Error(t, err)
+		})
+	}
+}
+
+func TestComputeHashCode(t *testing.T) {
+	expectedData := struct {
+		Field1 string
+		Field2 int
+	}{
+		Field1: "flytesnacks",
+		Field2: 42,
+	}
+	actualData := struct {
+		Field1 string
+		Field2 int
+	}{
+		Field1: "flytesnacks",
+		Field2: 42,
+	}
+	expectedHashcode, err := computeHashCode(expectedData)
+	assert.NoError(t, err)
+	actualHashcode, err := computeHashCode(actualData)
+	assert.NoError(t, err)
+	assert.Equal(t, expectedHashcode, actualHashcode)
+}