package br import ( "context" "errors" "testing" "time" v1alpha1 "github.com/pingcap/tidb-operator/api/v2/core/v1alpha1" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta" clusteragentv1 "github.com/tidbcloud/cluster-service-ng/api/tidbcloud/cluster-agent/v1" "github.com/tidbcloud/cluster-service-ng/config" "github.com/tidbcloud/cluster-service-ng/domain/entity" clustergatewaymocks "github.com/tidbcloud/cluster-service-ng/pkg/clustergateway/mocks" dataplanemocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplane/mocks" dbmocks "github.com/tidbcloud/cluster-service-ng/pkg/externalmocks/dbmocks" externalmocks "github.com/tidbcloud/cluster-service-ng/pkg/externalmocks/mocks" "github.com/tidbcloud/cluster-service-ng/pkg/storage/mysql/schema" "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo" repomocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks" "gorm.io/gorm" ) func createTestScaleManager(t *testing.T) (*scaleManagerImpl, *repomocks.SharedPoolRepo, *repomocks.ScaleTikvDecisionRepo, *clustergatewaymocks.Client) { t.Helper() sharedPoolRepo := repomocks.NewSharedPoolRepo(t) scaleTikvDecisionRepo := repomocks.NewScaleTikvDecisionRepo(t) dataplaneManager := dataplanemocks.NewManager(t) transactioner := dbmocks.NewTransactioner(t) clustergatewayCli := clustergatewaymocks.NewClient(t) bizConfig := &config.BizConfig{ DefaultTiBRConfig: config.TiBRConfig{ TiKVUsagePercentageThreshold: 80, }, } scaleManager := &scaleManagerImpl{ sharedPoolRepo: sharedPoolRepo, scaleTikvDecisionRepo: scaleTikvDecisionRepo, clustergatewayCli: clustergatewayCli, dataplaneManager: dataplaneManager, transactioner: transactioner, bizConfig: bizConfig, } return scaleManager, sharedPoolRepo, scaleTikvDecisionRepo, clustergatewayCli } func TestScaleManager_GetCurrentTiKVSpecUsage(t *testing.T) { tests := []struct { name string sharedPoolID string setupMocks func(*repomocks.SharedPoolRepo, *clustergatewaymocks.Client) expectedResult *schema.TiKVSpecUsage expectedError bool errorContains string }{ { name: "successful calculation with multiple nodes", sharedPoolID: "pool-1", setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, clustergatewayCli *clustergatewaymocks.Client) { sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-1", ComponentGroups: []*apiv1.Component{ {Id: "comp-1", ComponentType: apiv1.ComponentType_TIKV}, }, }, } sharedPoolRepo.On("Load", mock.Anything, "pool-1", false).Return(sharedPool, nil) topology := &clusteragentv1.GetClusterTopologyResp{ TikvNodes: []*clusteragentv1.GetClusterTopologyResp_TiKVNodeInfo{ { StorageCapacityBytes: uint64(100000), StorageAvailableBytes: uint64(30000), }, { StorageCapacityBytes: uint64(150000), StorageAvailableBytes: uint64(45000), }, { StorageCapacityBytes: uint64(200000), StorageAvailableBytes: uint64(60000), }, }, } clustergatewayCli.On("GetClusterTopology", mock.Anything, sharedPool, mock.MatchedBy(func(req *clusteragentv1.GetClusterTopologyReq) bool { return req.ClusterId == "pool-1" && req.IncludeTikvUsage })).Return(topology, nil) }, expectedResult: &schema.TiKVSpecUsage{ NodeCount: 3, TotalCapacityBytes: 450000, // 100000 + 150000 + 200000 TotalAvailableBytes: 135000, // 30000 + 45000 + 60000 UsagePercentage: 70, // ceil(100 - 135000 * 100 / 450000) = ceil(70) = 70 }, expectedError: false, }, { name: "successful calculation with single node", sharedPoolID: "pool-1", setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, clustergatewayCli *clustergatewaymocks.Client) { sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-1", ComponentGroups: []*apiv1.Component{ {Id: "comp-1", ComponentType: apiv1.ComponentType_TIKV}, }, }, } sharedPoolRepo.On("Load", mock.Anything, "pool-1", false).Return(sharedPool, nil) // Create a mock topology response - we'll fix the struct later topology := &clusteragentv1.GetClusterTopologyResp{ TikvNodes: []*clusteragentv1.GetClusterTopologyResp_TiKVNodeInfo{ { StorageCapacityBytes: uint64(100000), StorageAvailableBytes: uint64(75000), }, }, } clustergatewayCli.On("GetClusterTopology", mock.Anything, sharedPool, mock.MatchedBy(func(req *clusteragentv1.GetClusterTopologyReq) bool { return req.ClusterId == "pool-1" && req.IncludeTikvUsage })).Return(topology, nil) }, expectedResult: &schema.TiKVSpecUsage{ NodeCount: 1, TotalCapacityBytes: 100000, TotalAvailableBytes: 75000, UsagePercentage: 25, // ceil(100 - 75000 * 100 / 100000) = ceil(25) = 25 }, expectedError: false, }, { name: "zero available bytes", sharedPoolID: "pool-1", setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, clustergatewayCli *clustergatewaymocks.Client) { sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-1", ComponentGroups: []*apiv1.Component{ {Id: "comp-1", ComponentType: apiv1.ComponentType_TIKV}, }, }, } sharedPoolRepo.On("Load", mock.Anything, "pool-1", false).Return(sharedPool, nil) topology := &clusteragentv1.GetClusterTopologyResp{ TikvNodes: []*clusteragentv1.GetClusterTopologyResp_TiKVNodeInfo{ { StorageCapacityBytes: uint64(100000), StorageAvailableBytes: uint64(0), }, }, } clustergatewayCli.On("GetClusterTopology", mock.Anything, sharedPool, mock.MatchedBy(func(req *clusteragentv1.GetClusterTopologyReq) bool { return req.ClusterId == "pool-1" && req.IncludeTikvUsage })).Return(topology, nil) }, expectedResult: &schema.TiKVSpecUsage{ NodeCount: 1, TotalCapacityBytes: 100000, TotalAvailableBytes: 0, UsagePercentage: 100, }, expectedError: false, }, { name: "fractional percentage calculation", sharedPoolID: "pool-1", setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, clustergatewayCli *clustergatewaymocks.Client) { sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-1", ComponentGroups: []*apiv1.Component{ {Id: "comp-1", ComponentType: apiv1.ComponentType_TIKV}, }, }, } sharedPoolRepo.On("Load", mock.Anything, "pool-1", false).Return(sharedPool, nil) topology := &clusteragentv1.GetClusterTopologyResp{ TikvNodes: []*clusteragentv1.GetClusterTopologyResp_TiKVNodeInfo{ { StorageCapacityBytes: uint64(100000), StorageAvailableBytes: uint64(33333), }, }, } clustergatewayCli.On("GetClusterTopology", mock.Anything, sharedPool, mock.MatchedBy(func(req *clusteragentv1.GetClusterTopologyReq) bool { return req.ClusterId == "pool-1" && req.IncludeTikvUsage })).Return(topology, nil) }, expectedResult: &schema.TiKVSpecUsage{ NodeCount: 1, TotalCapacityBytes: 100000, TotalAvailableBytes: 33333, UsagePercentage: 67, // ceil(100 - 33333 * 100 / 100000) = ceil(66.667) = 67 }, expectedError: false, }, { name: "shared pool load error", sharedPoolID: "pool-1", setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, clustergatewayCli *clustergatewaymocks.Client) { sharedPoolRepo.On("Load", mock.Anything, "pool-1", false).Return(nil, assert.AnError) }, expectedResult: nil, expectedError: true, errorContains: "failed to load target sharedPool", }, { name: "cluster topology error", sharedPoolID: "pool-1", setupMocks: func(sharedPoolRepo *repomocks.SharedPoolRepo, clustergatewayCli *clustergatewaymocks.Client) { sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-1", ComponentGroups: []*apiv1.Component{ {Id: "comp-1", ComponentType: apiv1.ComponentType_TIKV}, }, }, } sharedPoolRepo.On("Load", mock.Anything, "pool-1", false).Return(sharedPool, nil) clustergatewayCli.On("GetClusterTopology", mock.Anything, sharedPool, mock.MatchedBy(func(req *clusteragentv1.GetClusterTopologyReq) bool { return req.ClusterId == "pool-1" && req.IncludeTikvUsage })).Return(nil, assert.AnError) }, expectedResult: nil, expectedError: true, errorContains: "failed to get cluster topology", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { scaleManager, sharedPoolRepo, _, clustergatewayCli := createTestScaleManager(t) tt.setupMocks(sharedPoolRepo, clustergatewayCli) result, err := scaleManager.getCurrentTiKVSpecUsage(context.Background(), tt.sharedPoolID) if tt.expectedError { require.Error(t, err) require.Contains(t, err.Error(), tt.errorContains) require.Nil(t, result) } else { require.NoError(t, err) require.NotNil(t, result) assert.Equal(t, tt.expectedResult.NodeCount, result.NodeCount) assert.Equal(t, tt.expectedResult.TotalCapacityBytes, result.TotalCapacityBytes) assert.Equal(t, tt.expectedResult.TotalAvailableBytes, result.TotalAvailableBytes) assert.Equal(t, tt.expectedResult.UsagePercentage, result.UsagePercentage) } }) } } func TestCalScaleDecisionTo(t *testing.T) { tests := []struct { name string toFullfilDataBytes uint64 basedUsage *schema.TiKVSpecUsage threshold uint32 expectedResult *schema.TiKVSpecUsage }{ { name: "no scaling needed - big number", toFullfilDataBytes: 1011324, basedUsage: &schema.TiKVSpecUsage{ NodeCount: 3, TotalCapacityBytes: 4752541396992, TotalAvailableBytes: 4750925886378, UsagePercentage: 1, }, threshold: 80, expectedResult: &schema.TiKVSpecUsage{ NodeCount: 3, TotalCapacityBytes: 4752541396992, TotalAvailableBytes: 4750924875054, UsagePercentage: 1, }, }, { name: "scaling needed - big number", toFullfilDataBytes: 4752541396992, basedUsage: &schema.TiKVSpecUsage{ NodeCount: 3, TotalCapacityBytes: 4752541396992, TotalAvailableBytes: 4750925886378, UsagePercentage: 1, }, threshold: 80, expectedResult: &schema.TiKVSpecUsage{ NodeCount: 6, // raw target is 4, then round up to the next multiple of 3 TotalCapacityBytes: 4752541396992 / 3 * 6, TotalAvailableBytes: 4750925886378, UsagePercentage: 51, }, }, { name: "no scaling needed - usage below threshold", toFullfilDataBytes: 100000, basedUsage: &schema.TiKVSpecUsage{ NodeCount: 3, TotalCapacityBytes: 1000000, TotalAvailableBytes: 800000, UsagePercentage: 20, }, threshold: 80, expectedResult: &schema.TiKVSpecUsage{ NodeCount: 3, TotalCapacityBytes: 1000000, TotalAvailableBytes: 700000, // 800000 - 100000 UsagePercentage: 30, // ceil((1000000 - 700000) * 100 / 1000000) }, }, { name: "scaling needed - usage exceeds threshold", toFullfilDataBytes: 500000, basedUsage: &schema.TiKVSpecUsage{ NodeCount: 2, TotalCapacityBytes: 1000000, TotalAvailableBytes: 200000, UsagePercentage: 80, }, threshold: 80, expectedResult: &schema.TiKVSpecUsage{ NodeCount: 6, // raw target is 4, then round up to the next multiple of 3 TotalCapacityBytes: 3000000, // 6 * 500000 TotalAvailableBytes: 1700000, // 3000000 - 1300000 UsagePercentage: 44, // ceil(1300000 * 100 / 3000000) }, }, { name: "edge case - zero available bytes", toFullfilDataBytes: 100000, basedUsage: &schema.TiKVSpecUsage{ NodeCount: 1, TotalCapacityBytes: 1000000, TotalAvailableBytes: 0, UsagePercentage: 100, }, threshold: 80, expectedResult: &schema.TiKVSpecUsage{ NodeCount: 3, // raw target is 2, then round up to the next multiple of 3 TotalCapacityBytes: 3000000, // 3 * 1000000 TotalAvailableBytes: 1900000, // 3000000 - 1100000 UsagePercentage: 37, // ceil(1100000 * 100 / 3000000) }, }, { name: "edge case - zero data to fulfill", toFullfilDataBytes: 0, basedUsage: &schema.TiKVSpecUsage{ NodeCount: 3, TotalCapacityBytes: 1000000, TotalAvailableBytes: 500000, UsagePercentage: 50, }, threshold: 80, expectedResult: &schema.TiKVSpecUsage{ NodeCount: 3, TotalCapacityBytes: 1000000, TotalAvailableBytes: 500000, UsagePercentage: 50, }, }, { name: "large data requirement", toFullfilDataBytes: 2000000, basedUsage: &schema.TiKVSpecUsage{ NodeCount: 2, TotalCapacityBytes: 1000000, TotalAvailableBytes: 100000, UsagePercentage: 90, }, threshold: 80, expectedResult: &schema.TiKVSpecUsage{ NodeCount: 9, // raw target is 8, then round up to the next multiple of 3 TotalCapacityBytes: 4500000, // 9 * 500000 TotalAvailableBytes: 1600000, // 4500000 - 2900000 UsagePercentage: 65, // ceil(2900000 * 100 / 4500000) }, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { result := calScaleDecisionTo(tt.toFullfilDataBytes, tt.basedUsage, tt.threshold) assert.Equal(t, tt.expectedResult.NodeCount, result.NodeCount, "NodeCount should match") assert.Equal(t, tt.expectedResult.TotalCapacityBytes, result.TotalCapacityBytes, "TotalCapacityBytes should match") assert.Equal(t, tt.expectedResult.TotalAvailableBytes, result.TotalAvailableBytes, "TotalAvailableBytes should match") assert.Equal(t, tt.expectedResult.UsagePercentage, result.UsagePercentage, "UsagePercentage should match") }) } } // Helper function to create test decisions func createTestDecision(id uint, baseOn string, baseID uint, createdAt time.Time) *schema.ScaleTikvDecision { decision := &schema.ScaleTikvDecision{ BaseOn: baseOn, BaseID: baseID, } decision.ID = id decision.CreatedAt = createdAt return decision } func TestBuildDecisionChains(t *testing.T) { t.Run("single decision chain - linear progression", func(t *testing.T) { decisions := []*schema.ScaleTikvDecision{ createTestDecision(1, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), createTestDecision(2, schema.BaseOnAnotherDecision, 1, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), createTestDecision(3, schema.BaseOnAnotherDecision, 2, time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC)), } result := buildDecisionChains(decisions) // Should have 1 chain assert.Len(t, result, 1, "Should have 1 chain") // Check the chain structure chain := result[0] assert.Equal(t, map[uint]bool{1: true, 2: true, 3: true}, chain.ids, "Chain should contain all decision IDs") assert.Len(t, chain.reversedOrder, 3, "Chain should have 3 decisions in reversed order") // Check the order (should be 3 -> 2 -> 1) assert.Equal(t, uint(3), chain.reversedOrder[0].ID, "First decision should be ID 3") assert.Equal(t, uint(2), chain.reversedOrder[1].ID, "Second decision should be ID 2") assert.Equal(t, uint(1), chain.reversedOrder[2].ID, "Third decision should be ID 1") }) t.Run("empty decisions list", func(t *testing.T) { decisions := []*schema.ScaleTikvDecision{} result := buildDecisionChains(decisions) assert.Empty(t, result, "Should have 0 chains") }) t.Run("single decision - no chain", func(t *testing.T) { decisions := []*schema.ScaleTikvDecision{ createTestDecision(1, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), } result := buildDecisionChains(decisions) // Should have 1 chain with just the single decision assert.Len(t, result, 1, "Should have 1 chain") chain := result[0] assert.Equal(t, map[uint]bool{1: true}, chain.ids, "Chain should contain decision ID 1") assert.Len(t, chain.reversedOrder, 1, "Chain should have 1 decision") assert.Equal(t, uint(1), chain.reversedOrder[0].ID, "Decision should be ID 1") }) t.Run("multiple decision chain", func(t *testing.T) { decisions := []*schema.ScaleTikvDecision{ createTestDecision(1, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), createTestDecision(2, schema.BaseOnAnotherDecision, 1, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), createTestDecision(3, schema.BaseOnAnotherDecision, 2, time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC)), createTestDecision(4, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), createTestDecision(5, schema.BaseOnAnotherDecision, 4, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), createTestDecision(6, schema.BaseOnAnotherDecision, 5, time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC)), } result := buildDecisionChains(decisions) // Should have 1 chain assert.Len(t, result, 2, "Should have 1 chain") // Check the chain structure chain1 := result[0] assert.Equal(t, map[uint]bool{1: true, 2: true, 3: true}, chain1.ids, "Chain should contain all decision IDs") assert.Len(t, chain1.reversedOrder, 3, "Chain should have 3 decisions in reversed order") // Check the order (should be 3 -> 2 -> 1) assert.Equal(t, uint(3), chain1.reversedOrder[0].ID, "First decision should be ID 3") assert.Equal(t, uint(2), chain1.reversedOrder[1].ID, "Second decision should be ID 2") assert.Equal(t, uint(1), chain1.reversedOrder[2].ID, "Third decision should be ID 1") chain2 := result[1] assert.Equal(t, map[uint]bool{4: true, 5: true, 6: true}, chain2.ids, "Chain should contain all decision IDs") assert.Len(t, chain2.reversedOrder, 3, "Chain should have 3 decisions in reversed order") assert.Equal(t, uint(6), chain2.reversedOrder[0].ID, "First decision should be ID 6") assert.Equal(t, uint(5), chain2.reversedOrder[1].ID, "Second decision should be ID 5") assert.Equal(t, uint(4), chain2.reversedOrder[2].ID, "Third decision should be ID 4") }) t.Run("base is misisng in decision chain", func(t *testing.T) { decisions := []*schema.ScaleTikvDecision{ createTestDecision(2, schema.BaseOnAnotherDecision, 1, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), createTestDecision(3, schema.BaseOnAnotherDecision, 2, time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC)), } result := buildDecisionChains(decisions) // Should have 1 chain assert.Len(t, result, 1, "Should have 1 chain") // Check the chain structure chain := result[0] assert.Equal(t, map[uint]bool{2: true, 3: true}, chain.ids, "Chain should contain all decision IDs") assert.Len(t, chain.reversedOrder, 2, "Chain should have 2 decisions in reversed order") // Check the order (should be 3 -> 2 -> 1) assert.Equal(t, uint(3), chain.reversedOrder[0].ID, "First decision should be ID 3") assert.Equal(t, uint(2), chain.reversedOrder[1].ID, "Second decision should be ID 2") }) } func TestAssembleScaleDecision(t *testing.T) { tests := []struct { name string scaleInput *ScaleInput scaleFrom *schema.TiKVSpecUsage scaleTo *schema.TiKVSpecUsage baseOn string baseId uint expectedResult *schema.ScaleTikvDecision }{ { name: "scaling up - need scale", scaleInput: &ScaleInput{ SharedPoolID: "pool-1", BizType: "restore", BizID: "restore-123", }, scaleFrom: &schema.TiKVSpecUsage{ NodeCount: 3, TotalCapacityBytes: 300000, TotalAvailableBytes: 100000, UsagePercentage: 67, }, scaleTo: &schema.TiKVSpecUsage{ NodeCount: 5, TotalCapacityBytes: 500000, TotalAvailableBytes: 200000, UsagePercentage: 60, }, baseOn: schema.BaseOnTiKVUsage, baseId: 0, expectedResult: &schema.ScaleTikvDecision{ SharedPoolID: "pool-1", BizType: "restore", BizID: "restore-123", Status: schema.ScaleStatusExecuting, NeedScale: true, BaseOn: schema.BaseOnTiKVUsage, BaseID: 0, From: schema.TiKVSpecUsage{ NodeCount: 3, TotalCapacityBytes: 300000, TotalAvailableBytes: 100000, UsagePercentage: 67, }, To: schema.TiKVSpecUsage{ NodeCount: 5, TotalCapacityBytes: 500000, TotalAvailableBytes: 200000, UsagePercentage: 60, }, }, }, { name: "same node count - no scale needed", scaleInput: &ScaleInput{ SharedPoolID: "pool-3", BizType: "migration", BizID: "migration-789", }, scaleFrom: &schema.TiKVSpecUsage{ NodeCount: 4, TotalCapacityBytes: 400000, TotalAvailableBytes: 150000, UsagePercentage: 63, }, scaleTo: &schema.TiKVSpecUsage{ NodeCount: 4, TotalCapacityBytes: 400000, TotalAvailableBytes: 150000, UsagePercentage: 63, }, baseOn: schema.BaseOnAnotherDecision, baseId: 456, expectedResult: &schema.ScaleTikvDecision{ SharedPoolID: "pool-3", BizType: "migration", BizID: "migration-789", Status: schema.ScaleStatusReachExpected, NeedScale: false, BaseOn: schema.BaseOnAnotherDecision, BaseID: 456, From: schema.TiKVSpecUsage{ NodeCount: 4, TotalCapacityBytes: 400000, TotalAvailableBytes: 150000, UsagePercentage: 63, }, To: schema.TiKVSpecUsage{ NodeCount: 4, TotalCapacityBytes: 400000, TotalAvailableBytes: 150000, UsagePercentage: 63, }, }, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { result := assembleScaleDecision(tt.scaleInput, tt.scaleFrom, tt.scaleTo, tt.baseOn, tt.baseId) assert.NotNil(t, result, "Result should not be nil") assert.Equal(t, tt.expectedResult.SharedPoolID, result.SharedPoolID, "SharedPoolID should match") assert.Equal(t, tt.expectedResult.BizType, result.BizType, "BizType should match") assert.Equal(t, tt.expectedResult.BizID, result.BizID, "BizID should match") assert.Equal(t, tt.expectedResult.NeedScale, result.NeedScale, "NeedScale should match") assert.Equal(t, tt.expectedResult.BaseOn, result.BaseOn, "BaseOn should match") assert.Equal(t, tt.expectedResult.BaseID, result.BaseID, "BaseID should match") // Test the core logic: status should be "executing" if needScale is true, "reach_expected" if false if tt.expectedResult.NeedScale { assert.Equal(t, "executing", result.Status, "Status should be executing when scaling is needed") } else { assert.Equal(t, "reach_expected", result.Status, "Status should be reach_expected when no scaling is needed") } // Test From field assert.Equal(t, tt.expectedResult.From.NodeCount, result.From.NodeCount, "From.NodeCount should match") assert.Equal(t, tt.expectedResult.From.TotalCapacityBytes, result.From.TotalCapacityBytes, "From.TotalCapacityBytes should match") assert.Equal(t, tt.expectedResult.From.TotalAvailableBytes, result.From.TotalAvailableBytes, "From.TotalAvailableBytes should match") assert.Equal(t, tt.expectedResult.From.UsagePercentage, result.From.UsagePercentage, "From.UsagePercentage should match") // Test To field assert.Equal(t, tt.expectedResult.To.NodeCount, result.To.NodeCount, "To.NodeCount should match") assert.Equal(t, tt.expectedResult.To.TotalCapacityBytes, result.To.TotalCapacityBytes, "To.TotalCapacityBytes should match") assert.Equal(t, tt.expectedResult.To.TotalAvailableBytes, result.To.TotalAvailableBytes, "To.TotalAvailableBytes should match") assert.Equal(t, tt.expectedResult.To.UsagePercentage, result.To.UsagePercentage, "To.UsagePercentage should match") }) } } func TestGetBasedDecision(t *testing.T) { tests := []struct { name string sharedPoolID string now time.Time setupMocks func(*repomocks.ScaleTikvDecisionRepo) expectedResult *schema.ScaleTikvDecision expectedError bool errorContains string }{ { name: "no executed fulfilling decisions", sharedPoolID: "pool-1", now: time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC), setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo) { // First call: get executed fulfilling decisions (should return empty) twoDaysAgo := time.Date(2023, 12, 30, 12, 0, 0, 0, time.UTC) scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: false, StartTime: &twoDaysAgo, }).Return([]*schema.ScaleTikvDecision{}, nil) }, expectedResult: nil, expectedError: false, }, { name: "single decision chain - return latest decision", sharedPoolID: "pool-1", now: time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC), setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo) { // First call: get executed fulfilling decisions fulfillingDecisions := []*schema.ScaleTikvDecision{ createTestDecision(3, schema.BaseOnAnotherDecision, 2, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), } twoDaysAgo := time.Date(2023, 12, 30, 12, 0, 0, 0, time.UTC) scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: false, StartTime: &twoDaysAgo, }).Return(fulfillingDecisions, nil) // Second call: get all decisions in the last 2 days allDecisions := []*schema.ScaleTikvDecision{ createTestDecision(1, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 9, 0, 0, 0, time.UTC)), createTestDecision(2, schema.BaseOnAnotherDecision, 1, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), createTestDecision(3, schema.BaseOnAnotherDecision, 2, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), createTestDecision(4, schema.BaseOnAnotherDecision, 3, time.Date(2024, 1, 1, 11, 0, 0, 1, time.UTC)), } scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: true, StartTime: &twoDaysAgo, }).Return(allDecisions, nil) }, expectedResult: createTestDecision(4, schema.BaseOnAnotherDecision, 3, time.Date(2024, 1, 1, 11, 0, 0, 1, time.UTC)), expectedError: false, }, { name: "multiple decision chains - return from correct chain", sharedPoolID: "pool-1", now: time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC), setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo) { // First call: get executed fulfilling decisions fulfillingDecisions := []*schema.ScaleTikvDecision{ createTestDecision(5, schema.BaseOnAnotherDecision, 4, time.Date(2024, 1, 1, 11, 30, 0, 0, time.UTC)), } twoDaysAgo := time.Date(2023, 12, 30, 12, 0, 0, 0, time.UTC) scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: false, StartTime: &twoDaysAgo, }).Return(fulfillingDecisions, nil) // Second call: get all decisions in the last 2 days allDecisions := []*schema.ScaleTikvDecision{ // Chain 1: 1 -> 2 -> 3 createTestDecision(1, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 9, 0, 0, 0, time.UTC)), createTestDecision(2, schema.BaseOnAnotherDecision, 1, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), createTestDecision(3, schema.BaseOnAnotherDecision, 2, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), // Chain 2: 4 -> 5 createTestDecision(4, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 9, 30, 0, 0, time.UTC)), createTestDecision(5, schema.BaseOnAnotherDecision, 4, time.Date(2024, 1, 1, 11, 30, 0, 0, time.UTC)), } scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: true, StartTime: &twoDaysAgo, }).Return(allDecisions, nil) }, expectedResult: createTestDecision(5, schema.BaseOnAnotherDecision, 4, time.Date(2024, 1, 1, 11, 30, 0, 0, time.UTC)), expectedError: false, }, { name: "latest fulfilling decision not in any chain", sharedPoolID: "pool-1", now: time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC), setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo) { // First call: get executed fulfilling decisions fulfillingDecisions := []*schema.ScaleTikvDecision{ createTestDecision(3, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), } twoDaysAgo := time.Date(2023, 12, 30, 12, 0, 0, 0, time.UTC) scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: false, StartTime: &twoDaysAgo, }).Return(fulfillingDecisions, nil) // Second call: get all decisions in the last 2 days allDecisions := []*schema.ScaleTikvDecision{ createTestDecision(1, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 9, 0, 0, 0, time.UTC)), createTestDecision(2, schema.BaseOnAnotherDecision, 1, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), createTestDecision(3, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), } scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: true, StartTime: &twoDaysAgo, }).Return(allDecisions, nil) }, expectedResult: createTestDecision(3, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), expectedError: false, }, { name: "error listing fulfilling decisions", sharedPoolID: "pool-1", now: time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC), setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo) { twoDaysAgo := time.Date(2023, 12, 30, 12, 0, 0, 0, time.UTC) scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: false, StartTime: &twoDaysAgo, }).Return(nil, assert.AnError) }, expectedResult: nil, expectedError: true, errorContains: "failed to list other running scale decisions", }, { name: "error listing all decisions", sharedPoolID: "pool-1", now: time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC), setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo) { // First call: get executed fulfilling decisions fulfillingDecisions := []*schema.ScaleTikvDecision{ createTestDecision(1, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), } twoDaysAgo := time.Date(2023, 12, 30, 12, 0, 0, 0, time.UTC) scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: false, StartTime: &twoDaysAgo, }).Return(fulfillingDecisions, nil) // Second call: get all decisions in the last 2 days - should fail scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: true, StartTime: &twoDaysAgo, }).Return(nil, assert.AnError) }, expectedResult: nil, expectedError: true, errorContains: "failed to list two days scale decisions", }, { name: "complex chain with missing base decisions", sharedPoolID: "pool-1", now: time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC), setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo) { // First call: get executed fulfilling decisions fulfillingDecisions := []*schema.ScaleTikvDecision{ createTestDecision(4, schema.BaseOnAnotherDecision, 3, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), } twoDaysAgo := time.Date(2023, 12, 30, 12, 0, 0, 0, time.UTC) scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: false, StartTime: &twoDaysAgo, }).Return(fulfillingDecisions, nil) // Second call: get all decisions in the last 2 days // Decision 1 is missing (outside 2-day window), but decision 2 references it allDecisions := []*schema.ScaleTikvDecision{ createTestDecision(2, schema.BaseOnAnotherDecision, 1, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), // base decision 1 is missing createTestDecision(3, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 10, 30, 0, 0, time.UTC)), createTestDecision(4, schema.BaseOnAnotherDecision, 3, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), } scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: true, StartTime: &twoDaysAgo, }).Return(allDecisions, nil) }, expectedResult: createTestDecision(4, schema.BaseOnAnotherDecision, 3, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), expectedError: false, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { scaleManager, _, scaleTikvDecisionRepo, _ := createTestScaleManager(t) tt.setupMocks(scaleTikvDecisionRepo) result, err := scaleManager.getBasedDecision(context.Background(), tt.sharedPoolID, tt.now) if tt.expectedError { require.Error(t, err) require.Contains(t, err.Error(), tt.errorContains) require.Nil(t, result) } else { require.NoError(t, err) if tt.expectedResult == nil { require.Nil(t, result) } else { require.NotNil(t, result) assert.Equal(t, tt.expectedResult.ID, result.ID) assert.Equal(t, tt.expectedResult.BaseOn, result.BaseOn) assert.Equal(t, tt.expectedResult.BaseID, result.BaseID) assert.Equal(t, tt.expectedResult.CreatedAt, result.CreatedAt) } } }) } } func TestGetLatestDecision(t *testing.T) { tests := []struct { name string decisions []*schema.ScaleTikvDecision expectedResult *schema.ScaleTikvDecision }{ { name: "single decision", decisions: []*schema.ScaleTikvDecision{ createTestDecision(1, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), }, expectedResult: createTestDecision(1, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), }, { name: "multiple decisions - latest at the end", decisions: []*schema.ScaleTikvDecision{ createTestDecision(1, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 9, 0, 0, 0, time.UTC)), createTestDecision(2, schema.BaseOnAnotherDecision, 1, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), createTestDecision(3, schema.BaseOnAnotherDecision, 2, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), }, expectedResult: createTestDecision(3, schema.BaseOnAnotherDecision, 2, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), }, { name: "multiple decisions - latest at the beginning", decisions: []*schema.ScaleTikvDecision{ createTestDecision(3, schema.BaseOnAnotherDecision, 2, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), createTestDecision(1, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 9, 0, 0, 0, time.UTC)), createTestDecision(2, schema.BaseOnAnotherDecision, 1, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), }, expectedResult: createTestDecision(3, schema.BaseOnAnotherDecision, 2, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), }, { name: "multiple decisions - latest in the middle", decisions: []*schema.ScaleTikvDecision{ createTestDecision(1, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 9, 0, 0, 0, time.UTC)), createTestDecision(3, schema.BaseOnAnotherDecision, 2, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), createTestDecision(2, schema.BaseOnAnotherDecision, 1, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), }, expectedResult: createTestDecision(3, schema.BaseOnAnotherDecision, 2, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { result := getLatestDecision(tt.decisions) require.NotNil(t, result, "Result should not be nil") assert.Equal(t, tt.expectedResult.ID, result.ID, "Decision ID should match") assert.Equal(t, tt.expectedResult.BaseOn, result.BaseOn, "BaseOn should match") assert.Equal(t, tt.expectedResult.BaseID, result.BaseID, "BaseID should match") assert.Equal(t, tt.expectedResult.CreatedAt, result.CreatedAt, "CreatedAt should match") }) } } func TestAnyFullfillingDecision(t *testing.T) { tests := []struct { name string sharedPoolID string setupMocks func(*repomocks.ScaleTikvDecisionRepo) expectedResult bool expectedError bool errorContains string }{ { name: "no fulfilling decisions", sharedPoolID: "pool-1", setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo) { scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: false, }).Return([]*schema.ScaleTikvDecision{}, nil) }, expectedResult: false, expectedError: false, }, { name: "single fulfilling decision", sharedPoolID: "pool-1", setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo) { decisions := []*schema.ScaleTikvDecision{ createTestDecision(1, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), } scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: false, }).Return(decisions, nil) }, expectedResult: true, expectedError: false, }, { name: "multiple fulfilling decisions", sharedPoolID: "pool-1", setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo) { decisions := []*schema.ScaleTikvDecision{ createTestDecision(1, schema.BaseOnTiKVUsage, 0, time.Date(2024, 1, 1, 10, 0, 0, 0, time.UTC)), createTestDecision(2, schema.BaseOnAnotherDecision, 1, time.Date(2024, 1, 1, 11, 0, 0, 0, time.UTC)), createTestDecision(3, schema.BaseOnAnotherDecision, 2, time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC)), } scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: false, }).Return(decisions, nil) }, expectedResult: true, expectedError: false, }, { name: "repository error", sharedPoolID: "pool-1", setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo) { scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: false, }).Return(nil, errors.New("database connection failed")) }, expectedResult: false, expectedError: true, errorContains: "failed to list scale decisions", }, { name: "nil decisions slice", sharedPoolID: "pool-1", setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo) { scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{ SharedPoolID: "pool-1", IncludeDeleted: false, }).Return(nil, nil) }, expectedResult: false, expectedError: false, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { // Create mocks scaleTikvDecisionRepo := repomocks.NewScaleTikvDecisionRepo(t) clustergatewayClient := clustergatewaymocks.NewClient(t) dataplaneManager := dataplanemocks.NewManager(t) dbTransactioner := dbmocks.NewTransactioner(t) // Setup mocks tt.setupMocks(scaleTikvDecisionRepo) // Create scale manager scaleManager := &scaleManagerImpl{ scaleTikvDecisionRepo: scaleTikvDecisionRepo, clustergatewayCli: clustergatewayClient, dataplaneManager: dataplaneManager, transactioner: dbTransactioner, } // Call the function result, err := scaleManager.AnyFullfillingDecision(context.Background(), tt.sharedPoolID) // Assert results if tt.expectedError { require.Error(t, err) if tt.errorContains != "" { assert.Contains(t, err.Error(), tt.errorContains) } assert.Equal(t, tt.expectedResult, result) } else { require.NoError(t, err) assert.Equal(t, tt.expectedResult, result) } // Verify all expectations were met scaleTikvDecisionRepo.AssertExpectations(t) }) } } func TestWaitAndSetScaleFinish(t *testing.T) { tests := []struct { name string sharedPoolID string scaleID uint setupMocks func(*repomocks.ScaleTikvDecisionRepo, *repomocks.SharedPoolRepo, *dataplanemocks.Manager, *externalmocks.Client) expectedDone bool expectedError bool errorContains string }{ { name: "scale decision already reach expected", sharedPoolID: "pool-1", scaleID: 1, setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager, dataplaneClient *externalmocks.Client) { decision := &schema.ScaleTikvDecision{ Model: gorm.Model{ID: 1}, Status: schema.ScaleStatusReachExpected, } scaleTikvDecisionRepo.On("Load", mock.Anything, uint(1), true).Return(decision, nil) }, expectedDone: true, expectedError: false, }, { name: "scale decision is deleted", sharedPoolID: "pool-1", scaleID: 1, setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager, dataplaneClient *externalmocks.Client) { decision := &schema.ScaleTikvDecision{ Model: gorm.Model{ ID: 1, DeletedAt: gorm.DeletedAt{ Time: time.Now(), Valid: true, }, }, } scaleTikvDecisionRepo.On("Load", mock.Anything, uint(1), true).Return(decision, nil) }, expectedDone: true, expectedError: false, }, { name: "tikv group ready replicas >= target replicas", sharedPoolID: "pool-1", scaleID: 1, setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager, dataplaneClient *externalmocks.Client) { decision := &schema.ScaleTikvDecision{ Model: gorm.Model{ID: 1}, Status: "pending", To: schema.TiKVSpecUsage{ NodeCount: 3, }, } scaleTikvDecisionRepo.On("Load", mock.Anything, uint(1), true).Return(decision, nil) sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-1", ComponentGroups: []*apiv1.Component{ { ComponentType: apiv1.ComponentType_PD, Id: "1", }, }, }, } sharedPoolRepo.On("Load", mock.Anything, "pool-1", false).Return(sharedPool, nil) dataplaneManager.On("GetClient", mock.Anything).Return(dataplaneClient, nil) dataplaneClient.On("Get", mock.Anything, mock.Anything, mock.Anything).Run( func(args mock.Arguments) { obj := args.Get(2).(*v1alpha1.TiKVGroup) obj.Status.ReadyReplicas = 3 }, ).Return(nil) // Mock Save call scaleTikvDecisionRepo.On("Save", mock.Anything, mock.Anything).Run( func(args mock.Arguments) { decision := args.Get(1).(*schema.ScaleTikvDecision) assert.Equal(t, schema.ScaleStatusReachExpected, decision.Status) }, ).Return(nil) }, expectedDone: true, expectedError: false, }, { name: "tikv group ready replicas < target replicas", sharedPoolID: "pool-1", scaleID: 1, setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager, dataplaneClient *externalmocks.Client) { decision := &schema.ScaleTikvDecision{ Model: gorm.Model{ID: 1}, Status: "pending", To: schema.TiKVSpecUsage{ NodeCount: 5, }, } scaleTikvDecisionRepo.On("Load", mock.Anything, uint(1), true).Return(decision, nil) sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-1", ComponentGroups: []*apiv1.Component{ { ComponentType: apiv1.ComponentType_PD, Id: "1", }, }, }, } sharedPoolRepo.On("Load", mock.Anything, "pool-1", false).Return(sharedPool, nil) dataplaneManager.On("GetClient", mock.Anything).Return(dataplaneClient, nil) dataplaneClient.On("Get", mock.Anything, mock.Anything, mock.Anything).Run( func(args mock.Arguments) { obj := args.Get(2).(*v1alpha1.TiKVGroup) obj.Status.ReadyReplicas = 3 }, ).Return(nil) }, expectedDone: false, expectedError: false, }, { name: "error loading scale decision", sharedPoolID: "pool-1", scaleID: 1, setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager, dataplaneClient *externalmocks.Client) { scaleTikvDecisionRepo.On("Load", mock.Anything, uint(1), true).Return(nil, errors.New("database error")) }, expectedDone: false, expectedError: true, errorContains: "failed to load scale decision", }, { name: "error loading shared pool", sharedPoolID: "pool-1", scaleID: 1, setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager, dataplaneClient *externalmocks.Client) { decision := &schema.ScaleTikvDecision{ Model: gorm.Model{ID: 1}, Status: "pending", } scaleTikvDecisionRepo.On("Load", mock.Anything, uint(1), true).Return(decision, nil) sharedPoolRepo.On("Load", mock.Anything, "pool-1", false).Return(nil, errors.New("shared pool not found")) }, expectedDone: false, expectedError: true, errorContains: "failed to load shared pool", }, { name: "error getting dataplane client", sharedPoolID: "pool-1", scaleID: 1, setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager, dataplaneClient *externalmocks.Client) { decision := &schema.ScaleTikvDecision{ Model: gorm.Model{ID: 1}, Status: "pending", } scaleTikvDecisionRepo.On("Load", mock.Anything, uint(1), true).Return(decision, nil) sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-1", }, } sharedPoolRepo.On("Load", mock.Anything, "pool-1", false).Return(sharedPool, nil) dataplaneManager.On("GetClient", mock.Anything).Return(nil, errors.New("dataplane client error")) }, expectedDone: false, expectedError: true, errorContains: "failed to get dataplane client", }, { name: "error getting TiKVGroup", sharedPoolID: "pool-1", scaleID: 1, setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager, dataplaneClient *externalmocks.Client) { decision := &schema.ScaleTikvDecision{ Model: gorm.Model{ID: 1}, Status: "pending", To: schema.TiKVSpecUsage{ NodeCount: 3, }, } scaleTikvDecisionRepo.On("Load", mock.Anything, uint(1), true).Return(decision, nil) sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-1", ComponentGroups: []*apiv1.Component{ { ComponentType: apiv1.ComponentType_PD, Id: "1", }, }, }, } sharedPoolRepo.On("Load", mock.Anything, "pool-1", false).Return(sharedPool, nil) dataplaneManager.On("GetClient", mock.Anything).Return(dataplaneClient, nil) dataplaneClient.On("Get", mock.Anything, mock.Anything, mock.Anything).Return(errors.New("TiKVGroup not found")) }, expectedDone: false, expectedError: true, errorContains: "failed to get TiKVGroup", }, { name: "error saving scale decision", sharedPoolID: "pool-1", scaleID: 1, setupMocks: func(scaleTikvDecisionRepo *repomocks.ScaleTikvDecisionRepo, sharedPoolRepo *repomocks.SharedPoolRepo, dataplaneManager *dataplanemocks.Manager, dataplaneClient *externalmocks.Client) { decision := &schema.ScaleTikvDecision{ Model: gorm.Model{ID: 1}, Status: "pending", To: schema.TiKVSpecUsage{ NodeCount: 3, }, } scaleTikvDecisionRepo.On("Load", mock.Anything, uint(1), true).Return(decision, nil) sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-1", ComponentGroups: []*apiv1.Component{ { ComponentType: apiv1.ComponentType_PD, Id: "1", }, }, }, } sharedPoolRepo.On("Load", mock.Anything, "pool-1", false).Return(sharedPool, nil) dataplaneManager.On("GetClient", mock.Anything).Return(dataplaneClient, nil) dataplaneClient.On("Get", mock.Anything, mock.Anything, mock.Anything).Run( func(args mock.Arguments) { obj := args.Get(2).(*v1alpha1.TiKVGroup) obj.Status.ReadyReplicas = 3 }, ).Return(nil) // Mock Save call to return error scaleTikvDecisionRepo.On("Save", mock.Anything, mock.Anything).Run( func(args mock.Arguments) { decision := args.Get(1).(*schema.ScaleTikvDecision) assert.Equal(t, schema.ScaleStatusReachExpected, decision.Status) }, ).Return(errors.New("save error")) }, expectedDone: false, expectedError: true, errorContains: "failed to save scale decision", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { // Create mocks scaleTikvDecisionRepo := repomocks.NewScaleTikvDecisionRepo(t) sharedPoolRepo := repomocks.NewSharedPoolRepo(t) dataplaneManager := dataplanemocks.NewManager(t) dataplaneClient := externalmocks.NewClient(t) clustergatewayClient := clustergatewaymocks.NewClient(t) dbTransactioner := dbmocks.NewTransactioner(t) // Setup mocks tt.setupMocks(scaleTikvDecisionRepo, sharedPoolRepo, dataplaneManager, dataplaneClient) // Create scale manager scaleManager := &scaleManagerImpl{ scaleTikvDecisionRepo: scaleTikvDecisionRepo, sharedPoolRepo: sharedPoolRepo, dataplaneManager: dataplaneManager, clustergatewayCli: clustergatewayClient, transactioner: dbTransactioner, } // Call the function done, err := scaleManager.WaitAndSetScaleFinish(context.Background(), tt.sharedPoolID, tt.scaleID) // Assert results if tt.expectedError { require.Error(t, err) if tt.errorContains != "" { assert.Contains(t, err.Error(), tt.errorContains) } assert.Equal(t, tt.expectedDone, done) } else { require.NoError(t, err) assert.Equal(t, tt.expectedDone, done) } // Verify all expectations were met scaleTikvDecisionRepo.AssertExpectations(t) sharedPoolRepo.AssertExpectations(t) dataplaneManager.AssertExpectations(t) dataplaneClient.AssertExpectations(t) }) } }