package server import ( "context" "testing" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta" dplanemgrv1 "github.com/tidbcloud/cluster-service-ng/api/dataplane-manager/v1" "github.com/tidbcloud/cluster-service-ng/domain/entity" "github.com/tidbcloud/cluster-service-ng/domain/workflow" bractionmocks "github.com/tidbcloud/cluster-service-ng/domain/workflow/action/br/mocks" "github.com/tidbcloud/cluster-service-ng/domain/workflow/dag" clustergatewaymocks "github.com/tidbcloud/cluster-service-ng/pkg/clustergateway/mocks" dataplanemocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplanemgr/mocks" "github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions" "github.com/tidbcloud/cluster-service-ng/pkg/runtime/executor" executormocks "github.com/tidbcloud/cluster-service-ng/pkg/runtime/executor/mocks" repomnocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" "google.golang.org/protobuf/encoding/protojson" "google.golang.org/protobuf/proto" ) // Helper function to setup DAG mocks for mergeSharedPoolStatusFromDagTask func setupSharedPoolDAGMocks(dagExecutor *executormocks.RemoteExecutor) { // Mock DAG queries for mergeSharedPoolStatusFromDagTask // These calls use string planName and a function matcher dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanCreateSharedPool, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanDeleteSharedPool, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanUpdateComponent, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanApplyComponent, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanCreateComponent, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanDeleteTiKVGroup, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanSetSharedPoolPDDeploymentModeToMonolithic, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanSetSharedPoolPDDeploymentModeToMicroservice, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanCreateTiDBCluster, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanUpdateTiDBCluster, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanDeleteTiDBCluster, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanCreateTiDBGroup, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanUpdateTiDBGroup, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanDeleteTiDBGroup, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanEnableCoprocessorWorker, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanDisableCoprocessorWorker, mock.AnythingOfType("executor.MatchFilter")).Return(executor.NewExecTaskIdStatus(), nil) } func TestCreateTiDBGroup(t *testing.T) { tests := []struct { name string tidbGroup *entity.TiDBGroup setupMocks func(*repomnocks.TiDBGroupRepo, *repomnocks.TiDBClusterRepo, *repomnocks.SharedPoolRepo, *repomnocks.LockRepo, *repomnocks.RuConfigRepo, *executormocks.RemoteExecutor, *bractionmocks.ScaleManager) expectedError string expectedCode codes.Code validate func(*testing.T, *entity.TiDBGroup, error) }{ { name: "missing tidb group id", tidbGroup: &entity.TiDBGroup{ TiDBGroup: apiv1.TiDBGroup{ ClusterId: "testcluster", }, }, setupMocks: func(*repomnocks.TiDBGroupRepo, *repomnocks.TiDBClusterRepo, *repomnocks.SharedPoolRepo, *repomnocks.LockRepo, *repomnocks.RuConfigRepo, *executormocks.RemoteExecutor, *bractionmocks.ScaleManager) { }, expectedError: "tidb_group_id cannot be empty", expectedCode: codes.InvalidArgument, }, { name: "missing cluster id", tidbGroup: &entity.TiDBGroup{ TiDBGroup: apiv1.TiDBGroup{ Id: "testgroup", TenantId: "testtenant", }, }, setupMocks: func(*repomnocks.TiDBGroupRepo, *repomnocks.TiDBClusterRepo, *repomnocks.SharedPoolRepo, *repomnocks.LockRepo, *repomnocks.RuConfigRepo, *executormocks.RemoteExecutor, *bractionmocks.ScaleManager) { }, expectedError: "tidbGroup cluster id is required", expectedCode: codes.InvalidArgument, }, { name: "tidb group already exists", tidbGroup: &entity.TiDBGroup{ TiDBGroup: apiv1.TiDBGroup{ Id: "existinggroup", ClusterId: "testcluster", TenantId: "testtenant", }, }, setupMocks: func(tidbGroupRepo *repomnocks.TiDBGroupRepo, tidbClusterRepo *repomnocks.TiDBClusterRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, ruConfigRepo *repomnocks.RuConfigRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) { // Mock running DAG check (no running tasks) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanCreateTiDBGroup, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil) // Mock existing tidb group existingGroup := &entity.TiDBGroup{ TiDBGroup: apiv1.TiDBGroup{ Id: "existinggroup", ClusterId: "testcluster", Phase: apiv1.TiDBGroup_AVAILABLE, }, } tidbGroupRepo.On("Load", mock.Anything, "existinggroup", true).Return(existingGroup, nil) }, expectedError: "tidb group id already exists", expectedCode: codes.AlreadyExists, }, { name: "cluster not found", tidbGroup: &entity.TiDBGroup{ TiDBGroup: apiv1.TiDBGroup{ Id: "testgroup", ClusterId: "nonexistentcluster", TenantId: "testtenant", }, }, setupMocks: func(tidbGroupRepo *repomnocks.TiDBGroupRepo, tidbClusterRepo *repomnocks.TiDBClusterRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, ruConfigRepo *repomnocks.RuConfigRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) { // Mock running DAG check (no running tasks) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanCreateTiDBGroup, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil) // Mock tidb group not found tidbGroupRepo.On("Load", mock.Anything, "testgroup", true).Return(nil, status.Error(codes.NotFound, "not found")) // Mock lock operations lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil) // Mock cluster not found tidbClusterRepo.On("Load", mock.Anything, "nonexistentcluster", false).Return(nil, status.Error(codes.NotFound, "cluster not found")) }, expectedError: "cluster not found", expectedCode: codes.NotFound, }, { name: "valid dns identity format and preserve disable public dns identity", tidbGroup: &entity.TiDBGroup{ TiDBGroup: apiv1.TiDBGroup{ Id: "testgroup", ClusterId: "testcluster", TenantId: "testtenant", SqlEndpoint: &entity.SQLEndpoint{ DnsIdentity: "valid-dns-123", DisablePublicDnsIdentity: true, }, }, }, setupMocks: func(tidbGroupRepo *repomnocks.TiDBGroupRepo, tidbClusterRepo *repomnocks.TiDBClusterRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, ruConfigRepo *repomnocks.RuConfigRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) { // Mock running DAG check (no running tasks) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanCreateTiDBGroup, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil) // Setup shared pool DAG mocks setupSharedPoolDAGMocks(dagExecutor) // Mock tidb group not found tidbGroupRepo.On("Load", mock.Anything, "testgroup", true).Return(nil, status.Error(codes.NotFound, "not found")) // Mock lock operations lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil) // Mock cluster found cluster := &entity.TiDBCluster{ Id: "testcluster", SharedPoolId: "testsharedpool", Phase: apiv1.Cluster_AVAILABLE, Version: "v7.5.0", RequestUnit: &apiv1.Cluster_RequestUnit{ Min: 10, Max: 100, }, } tidbClusterRepo.On("Load", mock.Anything, "testcluster", false).Return(cluster, nil) // Mock shared pool found sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "testsharedpool", Phase: apiv1.SharedPool_AVAILABLE, DataplaneName: "test-dataplane", Provider: &apiv1.Provider{ Type: apiv1.ProviderType_AWS, }, ComponentGroups: []*apiv1.Component{ { Id: "default", ComponentType: apiv1.ComponentType_TIKV, Phase: apiv1.Component_AVAILABLE, }, }, }, EnableCapacityPlanner: true, } sharedPoolRepo.On("Load", mock.Anything, "testsharedpool", false).Return(sharedPool, nil) // Mock RU config ruConfig := &entity.RuConfig{ RuConfig: apiv1.RuConfig{ Components: map[string]*apiv1.RuConfig_ComponentsMap{ "tidb": { NodeSpecRuRange: map[string]*apiv1.RuConfig_RuRange{ "4c16g": { Min: 0, Max: 100000000, }, }, }, }, FeatureFlags: &apiv1.RuConfig_FeatureFlags{ EnableScaleSharedPoolByRu: proto.Bool(false), }, }, } ruConfigRepo.On("GetRuConfig", mock.Anything, "testsharedpool", apiv1.RuConfig_SHARED_POOL).Return(ruConfig, nil) // Mock DAG execution submission tidbGroupRepo.On("Save", mock.Anything, mock.Anything).Return(nil) dagExecutor.On("Submit", mock.Anything, dag.PlanCreateTiDBGroup, mock.MatchedBy(func(args actions.ActionArgs) bool { templateArgs, err := args.GetString(workflow.AutoscalerTemplateArgsKey) return err == nil && templateArgs != "" && args[workflow.AutoscalerTemplateArgsKey].String() == templateArgs }), mock.Anything).Return(nil) }, validate: func(t *testing.T, result *entity.TiDBGroup, err error) { t.Helper() require.NoError(t, err) require.NotNil(t, result) require.Equal(t, "valid-dns-123", result.GetSqlEndpoint().GetDnsIdentity()) require.True(t, result.GetSqlEndpoint().GetDisablePublicDnsIdentity()) }, }, { name: "empty dns identity should be auto-generated", tidbGroup: &entity.TiDBGroup{ TiDBGroup: apiv1.TiDBGroup{ Id: "testgroup", ClusterId: "testcluster", TenantId: "testtenant", SqlEndpoint: &entity.SQLEndpoint{ DnsIdentity: "", // Empty DNS identity }, }, }, setupMocks: func(tidbGroupRepo *repomnocks.TiDBGroupRepo, tidbClusterRepo *repomnocks.TiDBClusterRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, ruConfigRepo *repomnocks.RuConfigRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) { // Mock running DAG check (no running tasks) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanCreateTiDBGroup, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil) // Setup shared pool DAG mocks setupSharedPoolDAGMocks(dagExecutor) // Mock tidb group not found tidbGroupRepo.On("Load", mock.Anything, "testgroup", true).Return(nil, status.Error(codes.NotFound, "not found")) // Mock lock operations lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil) // Mock cluster found cluster := &entity.TiDBCluster{ Id: "testcluster", SharedPoolId: "testsharedpool", Phase: apiv1.Cluster_AVAILABLE, Version: "v7.5.0", RequestUnit: &apiv1.Cluster_RequestUnit{ Min: 10, Max: 100, }, } tidbClusterRepo.On("Load", mock.Anything, "testcluster", false).Return(cluster, nil) // Mock shared pool found sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "testsharedpool", Phase: apiv1.SharedPool_AVAILABLE, DataplaneName: "test-dataplane", Provider: &apiv1.Provider{ Type: apiv1.ProviderType_AWS, }, ComponentGroups: []*apiv1.Component{ { Id: "default", ComponentType: apiv1.ComponentType_TIKV, Phase: apiv1.Component_AVAILABLE, }, }, }, } sharedPoolRepo.On("Load", mock.Anything, "testsharedpool", false).Return(sharedPool, nil) // Mock RU config ruConfig := &entity.RuConfig{ RuConfig: apiv1.RuConfig{ Components: map[string]*apiv1.RuConfig_ComponentsMap{ "tidb": { NodeSpecRuRange: map[string]*apiv1.RuConfig_RuRange{ "4c16g": { Min: 0, Max: 100000000, }, }, }, }, FeatureFlags: &apiv1.RuConfig_FeatureFlags{ EnableScaleSharedPoolByRu: proto.Bool(false), }, }, } ruConfigRepo.On("GetRuConfig", mock.Anything, "testsharedpool", apiv1.RuConfig_SHARED_POOL).Return(ruConfig, nil) // Mock DAG execution submission tidbGroupRepo.On("Save", mock.Anything, mock.Anything).Return(nil) dagExecutor.On("Submit", mock.Anything, dag.PlanCreateTiDBGroup, mock.Anything, mock.Anything).Return(nil) }, validate: func(t *testing.T, result *entity.TiDBGroup, err error) { t.Helper() require.NoError(t, err) require.NotNil(t, result) require.NotEmpty(t, result.GetSqlEndpoint().GetDnsIdentity(), "DNS Identity should be auto-generated when empty") require.Greater(t, len(result.GetSqlEndpoint().GetDnsIdentity()), 6, "Auto-generated DNS Identity should have reasonable length") }, }, { name: "nil sql endpoint should be auto-generated with dns identity", tidbGroup: &entity.TiDBGroup{ TiDBGroup: apiv1.TiDBGroup{ Id: "testgroup", ClusterId: "testcluster", TenantId: "testtenant", SqlEndpoint: nil, // Nil SQL endpoint }, }, setupMocks: func(tidbGroupRepo *repomnocks.TiDBGroupRepo, tidbClusterRepo *repomnocks.TiDBClusterRepo, sharedPoolRepo *repomnocks.SharedPoolRepo, lockRepo *repomnocks.LockRepo, ruConfigRepo *repomnocks.RuConfigRepo, dagExecutor *executormocks.RemoteExecutor, scaleManager *bractionmocks.ScaleManager) { // Mock running DAG check (no running tasks) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanCreateTiDBGroup, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil) // Setup shared pool DAG mocks setupSharedPoolDAGMocks(dagExecutor) // Mock tidb group not found tidbGroupRepo.On("Load", mock.Anything, "testgroup", true).Return(nil, status.Error(codes.NotFound, "not found")) // Mock lock operations lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil) // Mock cluster found cluster := &entity.TiDBCluster{ Id: "testcluster", SharedPoolId: "testsharedpool", Phase: apiv1.Cluster_AVAILABLE, Version: "v7.5.0", RequestUnit: &apiv1.Cluster_RequestUnit{ Min: 10, Max: 100, }, } tidbClusterRepo.On("Load", mock.Anything, "testcluster", false).Return(cluster, nil) // Mock shared pool found sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "testsharedpool", Phase: apiv1.SharedPool_AVAILABLE, DataplaneName: "test-dataplane", Provider: &apiv1.Provider{ Type: apiv1.ProviderType_AWS, }, ComponentGroups: []*apiv1.Component{ { Id: "default", ComponentType: apiv1.ComponentType_TIKV, Phase: apiv1.Component_AVAILABLE, }, }, }, } sharedPoolRepo.On("Load", mock.Anything, "testsharedpool", false).Return(sharedPool, nil) // Mock RU config ruConfig := &entity.RuConfig{ RuConfig: apiv1.RuConfig{ Components: map[string]*apiv1.RuConfig_ComponentsMap{ "tidb": { NodeSpecRuRange: map[string]*apiv1.RuConfig_RuRange{ "4c16g": { Min: 0, Max: 100000000, }, }, }, }, FeatureFlags: &apiv1.RuConfig_FeatureFlags{ EnableScaleSharedPoolByRu: proto.Bool(false), }, }, } ruConfigRepo.On("GetRuConfig", mock.Anything, "testsharedpool", apiv1.RuConfig_SHARED_POOL).Return(ruConfig, nil) // Mock DAG execution submission tidbGroupRepo.On("Save", mock.Anything, mock.Anything).Return(nil) dagExecutor.On("Submit", mock.Anything, dag.PlanCreateTiDBGroup, mock.Anything, mock.Anything).Return(nil) }, validate: func(t *testing.T, result *entity.TiDBGroup, err error) { t.Helper() require.NoError(t, err) require.NotNil(t, result) require.NotNil(t, result.GetSqlEndpoint(), "SQL Endpoint should be created when nil") require.NotEmpty(t, result.GetSqlEndpoint().GetDnsIdentity(), "DNS Identity should be auto-generated") require.Greater(t, len(result.GetSqlEndpoint().GetDnsIdentity()), 6, "Auto-generated DNS Identity should have reasonable length") }, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { // Create mocks tidbGroupRepo := repomnocks.NewTiDBGroupRepo(t) tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t) sharedPoolRepo := repomnocks.NewSharedPoolRepo(t) lockRepo := repomnocks.NewLockRepo(t) ruConfigRepo := repomnocks.NewRuConfigRepo(t) dagExecutor := executormocks.NewRemoteExecutor(t) dataplaneClient := dataplanemocks.NewClient(t) clusterGatewayClient := clustergatewaymocks.NewClient(t) scaleManager := bractionmocks.NewScaleManager(t) // Setup default mock for ScaleManager (used in mergeSharedPoolStatusFromRestore) // Use Maybe() to make it optional since not all test cases reach this code path scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil).Maybe() // Setup mocks tt.setupMocks(tidbGroupRepo, tidbClusterRepo, sharedPoolRepo, lockRepo, ruConfigRepo, dagExecutor, scaleManager) // Setup cluster gateway client mock for successful test cases if tt.expectedError == "" && tt.validate != nil { clusterGatewayClient.On("GetKeySpace", mock.Anything, mock.Anything, mock.Anything).Return(nil, nil) // Mock dataplane client for ListNodePools (called in genNodeSpecMap) nodePoolResp := &dplanemgrv1.ListNodePoolsResponse{ NodePools: []*dplanemgrv1.NodePool{ { CpuCore: 8, MemoryGi: 16, NodeType: dplanemgrv1.NodeType_Dedicated, AvailableCpuMilliCore: 7000, AvailableMemoryMi: 15000, WithLocalDisk: false, }, { CpuCore: 4, MemoryGi: 16, NodeType: dplanemgrv1.NodeType_Dedicated, AvailableCpuMilliCore: 3470, AvailableMemoryMi: 13208, WithLocalDisk: false, }, }, } dataplaneClient.On("ListNodePools", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.ListNodePoolsRequest) bool { return req.DataPlaneId == "test-dataplane" })).Return(nodePoolResp, nil) } // Create server implementation server := &serverImpl{ tidbGroupRepo: tidbGroupRepo, tidbClusterRepo: tidbClusterRepo, sharedPoolRepo: sharedPoolRepo, lockRepo: lockRepo, ruConfigRepo: ruConfigRepo, dagExecutor: dagExecutor, transactioner: &mockTransactioner{}, dataplaneClient: dataplaneClient, clustergatewayCli: clusterGatewayClient, scaleManager: scaleManager, } // Execute test result, err := server.CreateTiDBGroup(context.Background(), tt.tidbGroup) // Validate results switch { case tt.expectedError != "": require.Error(t, err) require.Contains(t, err.Error(), tt.expectedError) if tt.expectedCode != codes.OK { require.Equal(t, tt.expectedCode, status.Code(err)) } require.Nil(t, result) case tt.validate != nil: tt.validate(t, result, err) default: require.NoError(t, err) require.NotNil(t, result) } }) } } func TestCreateTiDBGroup_PersistsRequestPrivateLBCIDRsWithoutDataplaneMerge(t *testing.T) { tidbGroupRepo := repomnocks.NewTiDBGroupRepo(t) tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t) sharedPoolRepo := repomnocks.NewSharedPoolRepo(t) lockRepo := repomnocks.NewLockRepo(t) ruConfigRepo := repomnocks.NewRuConfigRepo(t) dagExecutor := executormocks.NewRemoteExecutor(t) dataplaneClient := dataplanemocks.NewClient(t) clusterGatewayClient := clustergatewaymocks.NewClient(t) scaleManager := bractionmocks.NewScaleManager(t) setupSharedPoolDAGMocks(dagExecutor) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanCreateTiDBGroup, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil) scaleManager.On("AnyFullfillingDecision", mock.Anything, mock.Anything).Return(false, nil).Maybe() clusterGatewayClient.On("GetKeySpace", mock.Anything, mock.Anything, mock.Anything).Return(nil, nil) request := &entity.TiDBGroup{ TiDBGroup: apiv1.TiDBGroup{ Id: "testgroup", ClusterId: "testcluster", TenantId: "testtenant", SqlEndpoint: &entity.SQLEndpoint{ PrivateLb: &entity.LoadBalancer{ Enabled: true, Port: 4000, IpAllowList: []*entity.IpAllowList{ {CidrBlock: "192.168.0.0/24"}, }, }, }, }, } cluster := &entity.TiDBCluster{ Id: "testcluster", SharedPoolId: "testsharedpool", Phase: apiv1.Cluster_AVAILABLE, Version: "v7.5.0", RequestUnit: &apiv1.Cluster_RequestUnit{ Min: 10, Max: 100, }, } sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "testsharedpool", Phase: apiv1.SharedPool_AVAILABLE, DataplaneName: "test-dataplane", Provider: &apiv1.Provider{ Type: apiv1.ProviderType_AWS, }, ComponentGroups: []*apiv1.Component{ { Id: "default", ComponentType: apiv1.ComponentType_TIKV, Phase: apiv1.Component_AVAILABLE, }, }, }, } ruConfig := &entity.RuConfig{ RuConfig: apiv1.RuConfig{ Components: map[string]*apiv1.RuConfig_ComponentsMap{ "tidb": { NodeSpecRuRange: map[string]*apiv1.RuConfig_RuRange{ "4c16g": { Min: 0, Max: 100000000, }, }, }, }, FeatureFlags: &apiv1.RuConfig_FeatureFlags{ EnableScaleSharedPoolByRu: proto.Bool(false), }, }, } nodePoolResp := &dplanemgrv1.ListNodePoolsResponse{ NodePools: []*dplanemgrv1.NodePool{ { CpuCore: 4, MemoryGi: 16, NodeType: dplanemgrv1.NodeType_Dedicated, AvailableCpuMilliCore: 3470, AvailableMemoryMi: 13208, WithLocalDisk: false, }, }, } tidbGroupRepo.On("Load", mock.Anything, "testgroup", true).Return(nil, status.Error(codes.NotFound, "not found")) lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil) tidbClusterRepo.On("Load", mock.Anything, "testcluster", false).Return(cluster, nil) sharedPoolRepo.On("Load", mock.Anything, "testsharedpool", false).Return(sharedPool, nil) ruConfigRepo.On("GetRuConfig", mock.Anything, "testsharedpool", apiv1.RuConfig_SHARED_POOL).Return(ruConfig, nil) dataplaneClient.On("ListNodePools", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.ListNodePoolsRequest) bool { return req.DataPlaneId == "test-dataplane" })).Return(nodePoolResp, nil).Once() tidbGroupRepo.On("Save", mock.Anything, mock.MatchedBy(func(saved *entity.TiDBGroup) bool { got := saved.GetSqlEndpoint().GetPrivateLb().GetIpAllowList() if len(got) != 1 { return false } return got[0].GetCidrBlock() == "192.168.0.0/24" })).Return(nil).Once() dagExecutor.On("Submit", mock.Anything, dag.PlanCreateTiDBGroup, mock.Anything, mock.Anything).Return(nil).Once() server := &serverImpl{ tidbGroupRepo: tidbGroupRepo, tidbClusterRepo: tidbClusterRepo, sharedPoolRepo: sharedPoolRepo, lockRepo: lockRepo, ruConfigRepo: ruConfigRepo, dagExecutor: dagExecutor, transactioner: &mockTransactioner{}, dataplaneClient: dataplaneClient, clustergatewayCli: clusterGatewayClient, scaleManager: scaleManager, } result, err := server.CreateTiDBGroup(context.Background(), request) require.NoError(t, err) require.NotNil(t, result) require.Len(t, result.GetSqlEndpoint().GetPrivateLb().GetIpAllowList(), 1) require.Equal(t, "192.168.0.0/24", result.GetSqlEndpoint().GetPrivateLb().GetIpAllowList()[0].GetCidrBlock()) } func TestCreateTiDBGroup_TopologyDrivenInheritsTiDBSpec(t *testing.T) { tidbGroupRepo := repomnocks.NewTiDBGroupRepo(t) tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t) sharedPoolRepo := repomnocks.NewSharedPoolRepo(t) lockRepo := repomnocks.NewLockRepo(t) ruConfigRepo := repomnocks.NewRuConfigRepo(t) dagExecutor := executormocks.NewRemoteExecutor(t) dataplaneClient := dataplanemocks.NewClient(t) clusterGatewayClient := clustergatewaymocks.NewClient(t) scaleManager := bractionmocks.NewScaleManager(t) tidbCluster := &entity.TiDBCluster{ Id: "testcluster", TenantId: "testtenant", SharedPoolId: "testsharedpool", Phase: apiv1.Cluster_AVAILABLE, CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN, TopologySpec: &apiv1.TopologySpec{Components: []*apiv1.Component{ { Id: "testcluster", ComponentType: apiv1.ComponentType_TIDB, NodeQuota: &apiv1.NodeQuota{ NodeCount: 2, NodeSpec: &apiv1.NodeSpec{Cpu: 8, MemoryGi: 16}, }, }, { Id: "testcluster", ComponentType: apiv1.ComponentType_TIPROXY, NodeQuota: &apiv1.NodeQuota{ NodeCount: 1, NodeSpec: &apiv1.NodeSpec{Cpu: 4, MemoryGi: 8}, }, }, }}, } sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{ Id: "testsharedpool", Phase: apiv1.SharedPool_AVAILABLE, DataplaneName: "test-dataplane", CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN, Provider: &apiv1.Provider{Type: apiv1.ProviderType_AWS}, ComponentGroups: []*apiv1.Component{ { Id: "default", ComponentType: apiv1.ComponentType_TIKV, Phase: apiv1.Component_AVAILABLE, }, }, }} rq := &entity.TiDBGroup{TiDBGroup: apiv1.TiDBGroup{Id: "testgroup", TenantId: "testtenant", ClusterId: "testcluster"}} rq.TopologySpec = tidbCluster.GetTopologySpec() setupSharedPoolDAGMocks(dagExecutor) scaleManager.On("AnyFullfillingDecision", mock.Anything, "testsharedpool").Return(false, nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanCreateTiDBGroup, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil) tidbGroupRepo.On("Load", mock.Anything, "testgroup", true).Return(nil, status.Error(codes.NotFound, "not found")) lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil) tidbClusterRepo.On("Load", mock.Anything, "testcluster", false).Return(tidbCluster, nil) sharedPoolRepo.On("Load", mock.Anything, "testsharedpool", false).Return(sharedPool, nil) clusterGatewayClient.On("GetKeySpace", mock.Anything, sharedPool, "testcluster").Return(nil, nil) dataplaneClient.On("ListNodePools", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListNodePoolsResponse{ NodePools: []*dplanemgrv1.NodePool{ { CpuCore: 8, MemoryGi: 16, AvailableCpuMilliCore: 7500, AvailableMemoryMi: 15000, NodeType: dplanemgrv1.NodeType_Dedicated, }, }, }, nil) tidbGroupRepo.On("Save", mock.Anything, mock.Anything).Return(nil) dagExecutor.On("Submit", mock.Anything, dag.PlanCreateTiDBGroup, mock.Anything, mock.Anything).Run(func(args mock.Arguments) { actionArgs := args.Get(2).(actions.ActionArgs) sharedPoolStr, err := actionArgs.GetString(workflow.SharedPoolKey) require.NoError(t, err) var savedSharedPool entity.SharedPool require.NoError(t, protojson.Unmarshal([]byte(sharedPoolStr), &savedSharedPool)) var inherited *entity.Component for _, comp := range savedSharedPool.GetComponentGroups() { if comp.GetComponentType() == apiv1.ComponentType_TIDB && comp.GetId() == "testgroup" { inherited = comp break } } require.NotNil(t, inherited) require.NotNil(t, inherited.GetNodeQuota()) require.NotNil(t, inherited.GetNodeQuota().GetNodeSpec()) require.Equal(t, uint32(2), inherited.GetNodeQuota().GetNodeCount()) require.Equal(t, uint32(8), inherited.GetNodeQuota().GetNodeSpec().GetCpu()) require.Equal(t, uint32(16), inherited.GetNodeQuota().GetNodeSpec().GetMemoryGi()) require.Equal(t, apiv1.Component_CREATING, inherited.GetPhase()) }).Return(nil) server := &serverImpl{ tidbGroupRepo: tidbGroupRepo, tidbClusterRepo: tidbClusterRepo, sharedPoolRepo: sharedPoolRepo, lockRepo: lockRepo, ruConfigRepo: ruConfigRepo, dagExecutor: dagExecutor, transactioner: &mockTransactioner{}, dataplaneClient: dataplaneClient, clustergatewayCli: clusterGatewayClient, scaleManager: scaleManager, } result, err := server.CreateTiDBGroup(context.Background(), rq) require.NoError(t, err) require.NotNil(t, result) if result.GetSqlEndpoint() == nil || result.GetSqlEndpoint().GetDnsIdentity() == "" { t.Fatalf("sql endpoint dns identity should be set, got: %v", result.GetSqlEndpoint()) } } func TestCreateTiDBGroup_TopologyDrivenFillsNodeSpecAvailabilityFromDataplane(t *testing.T) { tidbGroupRepo := repomnocks.NewTiDBGroupRepo(t) tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t) sharedPoolRepo := repomnocks.NewSharedPoolRepo(t) lockRepo := repomnocks.NewLockRepo(t) ruConfigRepo := repomnocks.NewRuConfigRepo(t) dagExecutor := executormocks.NewRemoteExecutor(t) dataplaneClient := dataplanemocks.NewClient(t) clusterGatewayClient := clustergatewaymocks.NewClient(t) scaleManager := bractionmocks.NewScaleManager(t) tidbCluster := &entity.TiDBCluster{ Id: "testcluster", TenantId: "testtenant", SharedPoolId: "testsharedpool", Phase: apiv1.Cluster_AVAILABLE, CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN, } sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{ Id: "testsharedpool", Phase: apiv1.SharedPool_AVAILABLE, DataplaneName: "test-dataplane", CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN, Provider: &apiv1.Provider{ Type: apiv1.ProviderType_AWS, }, ComponentGroups: []*apiv1.Component{ { Id: "default", ComponentType: apiv1.ComponentType_TIKV, Phase: apiv1.Component_AVAILABLE, }, }, }} rq := &entity.TiDBGroup{TiDBGroup: apiv1.TiDBGroup{Id: "testgroup", TenantId: "testtenant", ClusterId: "testcluster"}} rq.TopologySpec = &apiv1.TopologySpec{Components: []*apiv1.Component{{ ComponentType: apiv1.ComponentType_TIDB, NodeQuota: &apiv1.NodeQuota{ NodeCount: 1, NodeSpec: &apiv1.NodeSpec{ Cpu: 8, MemoryGi: 16, }, }, }}} setupSharedPoolDAGMocks(dagExecutor) scaleManager.On("AnyFullfillingDecision", mock.Anything, "testsharedpool").Return(false, nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanCreateTiDBGroup, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil) tidbGroupRepo.On("Load", mock.Anything, "testgroup", true).Return(nil, status.Error(codes.NotFound, "not found")) lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil) tidbClusterRepo.On("Load", mock.Anything, "testcluster", false).Return(tidbCluster, nil) sharedPoolRepo.On("Load", mock.Anything, "testsharedpool", false).Return(sharedPool, nil) clusterGatewayClient.On("GetKeySpace", mock.Anything, sharedPool, "testcluster").Return(nil, nil) dataplaneClient.On("ListNodePools", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListNodePoolsResponse{ NodePools: []*dplanemgrv1.NodePool{ { CpuCore: 8, MemoryGi: 16, AvailableCpuMilliCore: 7500, AvailableMemoryMi: 15000, NodeType: dplanemgrv1.NodeType_Dedicated, }, }, }, nil) tidbGroupRepo.On("Save", mock.Anything, mock.Anything).Return(nil) dagExecutor.On("Submit", mock.Anything, dag.PlanCreateTiDBGroup, mock.Anything, mock.Anything).Run(func(args mock.Arguments) { actionArgs := args.Get(2).(actions.ActionArgs) sharedPoolStr, err := actionArgs.GetString(workflow.SharedPoolKey) require.NoError(t, err) var savedSharedPool entity.SharedPool require.NoError(t, protojson.Unmarshal([]byte(sharedPoolStr), &savedSharedPool)) var inherited *entity.Component for _, comp := range savedSharedPool.GetComponentGroups() { if comp.GetComponentType() == apiv1.ComponentType_TIDB && comp.GetId() == "testgroup" { inherited = comp break } } require.NotNil(t, inherited) require.NotNil(t, inherited.GetNodeQuota()) require.NotNil(t, inherited.GetNodeQuota().GetNodeSpec()) require.Positive(t, inherited.GetNodeQuota().GetNodeSpec().GetAvailableMemoryMi()) require.Positive(t, inherited.GetNodeQuota().GetNodeSpec().GetAvailableCpu()) }).Return(nil) server := &serverImpl{ tidbGroupRepo: tidbGroupRepo, tidbClusterRepo: tidbClusterRepo, sharedPoolRepo: sharedPoolRepo, lockRepo: lockRepo, ruConfigRepo: ruConfigRepo, dagExecutor: dagExecutor, transactioner: &mockTransactioner{}, dataplaneClient: dataplaneClient, clustergatewayCli: clusterGatewayClient, scaleManager: scaleManager, } result, err := server.CreateTiDBGroup(context.Background(), rq) require.NoError(t, err) require.NotNil(t, result) } func TestCreateTiDBGroup_TopologyDrivenRequiresRequestTopologySpec(t *testing.T) { tidbGroupRepo := repomnocks.NewTiDBGroupRepo(t) tidbClusterRepo := repomnocks.NewTiDBClusterRepo(t) sharedPoolRepo := repomnocks.NewSharedPoolRepo(t) lockRepo := repomnocks.NewLockRepo(t) ruConfigRepo := repomnocks.NewRuConfigRepo(t) dagExecutor := executormocks.NewRemoteExecutor(t) dataplaneClient := dataplanemocks.NewClient(t) clusterGatewayClient := clustergatewaymocks.NewClient(t) scaleManager := bractionmocks.NewScaleManager(t) tidbCluster := &entity.TiDBCluster{ Id: "testcluster", TenantId: "testtenant", SharedPoolId: "testsharedpool", Phase: apiv1.Cluster_AVAILABLE, CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN, } sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{ Id: "testsharedpool", Phase: apiv1.SharedPool_AVAILABLE, DataplaneName: "test-dataplane", CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN, Provider: &apiv1.Provider{Type: apiv1.ProviderType_AWS}, ComponentGroups: []*apiv1.Component{ { Id: "default", ComponentType: apiv1.ComponentType_TIKV, Phase: apiv1.Component_AVAILABLE, }, }, }} rq := &entity.TiDBGroup{TiDBGroup: apiv1.TiDBGroup{Id: "testgroup", TenantId: "testtenant", ClusterId: "testcluster"}} setupSharedPoolDAGMocks(dagExecutor) scaleManager.On("AnyFullfillingDecision", mock.Anything, "testsharedpool").Return(false, nil) dagExecutor.On("QueryRunnings", mock.Anything, dag.PlanCreateTiDBGroup, mock.Anything).Return(executor.NewExecTaskIdStatus(), nil) tidbGroupRepo.On("Load", mock.Anything, "testgroup", true).Return(nil, status.Error(codes.NotFound, "not found")) lockRepo.On("Lock", mock.Anything, mock.Anything).Return(func() error { return nil }, nil) tidbClusterRepo.On("Load", mock.Anything, "testcluster", false).Return(tidbCluster, nil) sharedPoolRepo.On("Load", mock.Anything, "testsharedpool", false).Return(sharedPool, nil) clusterGatewayClient.On("GetKeySpace", mock.Anything, sharedPool, "testcluster").Return(nil, nil) dataplaneClient.On("ListNodePools", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListNodePoolsResponse{ NodePools: []*dplanemgrv1.NodePool{ { CpuCore: 8, MemoryGi: 16, AvailableCpuMilliCore: 7500, AvailableMemoryMi: 15000, NodeType: dplanemgrv1.NodeType_Dedicated, }, }, }, nil) server := &serverImpl{ tidbGroupRepo: tidbGroupRepo, tidbClusterRepo: tidbClusterRepo, sharedPoolRepo: sharedPoolRepo, lockRepo: lockRepo, ruConfigRepo: ruConfigRepo, dagExecutor: dagExecutor, transactioner: &mockTransactioner{}, dataplaneClient: dataplaneClient, clustergatewayCli: clusterGatewayClient, scaleManager: scaleManager, } result, err := server.CreateTiDBGroup(context.Background(), rq) require.Nil(t, result) require.Error(t, err) require.Equal(t, codes.InvalidArgument, status.Code(err)) require.Contains(t, err.Error(), "topology spec is required") } func TestGeneratedTiDBGroupKeepsExistingRCUComponent(t *testing.T) { storage := &apiv1.Component_Storage{ Name: apiv1.Component_Storage_DATA, Type: "gp3", SizeGi: 100, } sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{ CapacityMode: apiv1.CapacityMode_RCU_DRIVEN, ComponentGroups: []*apiv1.Component{ { Id: "tidb-group", ComponentType: apiv1.ComponentType_TIDB, Phase: apiv1.Component_AVAILABLE, NodeQuota: &apiv1.NodeQuota{ NodeCount: 1, MinNodeCount: 1, MaxNodeCount: 10, NodeSpec: &apiv1.NodeSpec{Cpu: 2, MemoryGi: 8}, }, Storages: []*apiv1.Component_Storage{storage}, }, }, }} tidbCluster := &entity.TiDBCluster{RequestUnit: &apiv1.Cluster_RequestUnit{Min: 100, Max: 100}} tidbGroup := &entity.TiDBGroup{TiDBGroup: apiv1.TiDBGroup{Id: "tidb-group"}} nodeSpecMap := map[string][]*entity.NodeSpec{ "4C16G": {{Cpu: 4, MemoryGi: 16, AvailableCpu: 4, AvailableMemoryMi: 16384}}, } ruConfig := &entity.RuConfig{RuConfig: apiv1.RuConfig{ Components: map[string]*apiv1.RuConfig_ComponentsMap{ "tidb": { MaxNodeCount: 10, NodeSpecRuRange: map[string]*apiv1.RuConfig_RuRange{ "4c16g": {Min: 100, Max: 0, DefaultNodeCount: uint32Ptr(3)}, }, }, }, }} err := (&serverImpl{}).generatedTiDBGroup(context.Background(), sharedPool, tidbCluster, tidbGroup, nodeSpecMap, ruConfig) require.NoError(t, err) require.Len(t, sharedPool.GetComponentGroups(), 1) tidbComp := sharedPool.GetComponentGroups()[0] require.Equal(t, apiv1.Component_AVAILABLE, tidbComp.GetPhase()) require.Equal(t, uint32(1), tidbComp.GetNodeQuota().GetNodeCount()) require.Equal(t, uint32(2), tidbComp.GetNodeQuota().GetNodeSpec().GetCpu()) require.Equal(t, uint32(8), tidbComp.GetNodeQuota().GetNodeSpec().GetMemoryGi()) require.Same(t, storage, tidbComp.GetStorages()[0]) } func TestGeneratedTiDBGroupKeepsExistingRCUComponentWhenQuotaUnchanged(t *testing.T) { storage := &apiv1.Component_Storage{Name: apiv1.Component_Storage_DATA, Type: "gp3", SizeGi: 100} nodeSpec := &apiv1.NodeSpec{Cpu: 4, MemoryGi: 16, AvailableCpu: 4, AvailableMemoryMi: 16384} sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{ CapacityMode: apiv1.CapacityMode_RCU_DRIVEN, ComponentGroups: []*apiv1.Component{{ Id: "tidb-group", ComponentType: apiv1.ComponentType_TIDB, Phase: apiv1.Component_AVAILABLE, NodeQuota: &apiv1.NodeQuota{NodeCount: 3, MinNodeCount: 1, MaxNodeCount: 10, NodeSpec: nodeSpec}, Storages: []*apiv1.Component_Storage{storage}, }}, }} tidbCluster := &entity.TiDBCluster{RequestUnit: &apiv1.Cluster_RequestUnit{Min: 100, Max: 100}} tidbGroup := &entity.TiDBGroup{TiDBGroup: apiv1.TiDBGroup{Id: "tidb-group"}} nodeSpecMap := map[string][]*entity.NodeSpec{ "4C16G": {{Cpu: 4, MemoryGi: 16, AvailableCpu: 4, AvailableMemoryMi: 16384}}, } ruConfig := &entity.RuConfig{RuConfig: apiv1.RuConfig{Components: map[string]*apiv1.RuConfig_ComponentsMap{ "tidb": {MaxNodeCount: 10, NodeSpecRuRange: map[string]*apiv1.RuConfig_RuRange{ "4c16g": {Min: 100, Max: 0, DefaultNodeCount: uint32Ptr(3)}, }}, }}} err := (&serverImpl{}).generatedTiDBGroup(context.Background(), sharedPool, tidbCluster, tidbGroup, nodeSpecMap, ruConfig) require.NoError(t, err) require.Len(t, sharedPool.GetComponentGroups(), 1) require.Equal(t, apiv1.Component_AVAILABLE, sharedPool.GetComponentGroups()[0].GetPhase()) require.Same(t, storage, sharedPool.GetComponentGroups()[0].GetStorages()[0]) require.Equal(t, nodeSpec, sharedPool.GetComponentGroups()[0].GetNodeQuota().GetNodeSpec()) } func TestRegenerateTiDBGroupForRUScaleUpSkipsMixedDimensionTierTransition(t *testing.T) { currentNodeSpec := &entity.NodeSpec{Cpu: 4, MemoryGi: 16} tidbComp := &entity.Component{ Id: "tidb-group", ComponentType: apiv1.ComponentType_TIDB, Phase: apiv1.Component_AVAILABLE, NodeQuota: &entity.NodeQuota{ NodeCount: 1, MinNodeCount: 1, MaxNodeCount: 10, NodeSpec: currentNodeSpec, }, } sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{ Id: "shared-pool", CapacityMode: apiv1.CapacityMode_RCU_DRIVEN, ComponentGroups: []*apiv1.Component{tidbComp}, }} tidbCluster := &entity.TiDBCluster{ Id: "tidb-cluster", RequestUnit: &entity.RequestUnit{Min: 200, Max: 200}, } tidbGroup := &entity.TiDBGroup{TiDBGroup: apiv1.TiDBGroup{Id: tidbComp.GetId()}} ruConfig := &entity.RuConfig{RuConfig: apiv1.RuConfig{ Components: map[string]*apiv1.RuConfig_ComponentsMap{ "tidb": { NodeSpecRuRange: map[string]*apiv1.RuConfig_RuRange{ "4c16g": {Min: 0, Max: 100}, "8c8g": {Min: 101, Max: 0}, }, }, }, }} updated, err := (&serverImpl{}).regenerateTiDBGroupForRUScaleUp( context.Background(), sharedPool, &entity.RequestUnit{Min: 100, Max: 100}, tidbCluster, tidbGroup, nil, ruConfig, ) require.NoError(t, err) require.False(t, updated) require.Equal(t, apiv1.Component_AVAILABLE, tidbComp.GetPhase()) require.Same(t, currentNodeSpec, tidbComp.GetNodeQuota().GetNodeSpec()) } func TestGeneratedTiDBGroupKeepsExistingTopologyComponent(t *testing.T) { existing := &entity.Component{ Id: "tidb-group", ComponentType: apiv1.ComponentType_TIDB, Phase: apiv1.Component_AVAILABLE, NodeQuota: &entity.NodeQuota{NodeCount: 2, NodeSpec: &entity.NodeSpec{Cpu: 4, MemoryGi: 16}}, } sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{ CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN, ComponentGroups: []*apiv1.Component{existing}, }} err := (&serverImpl{}).generatedTiDBGroup( context.Background(), sharedPool, &entity.TiDBCluster{}, &entity.TiDBGroup{TiDBGroup: apiv1.TiDBGroup{Id: "tidb-group"}}, nil, nil, ) require.NoError(t, err) require.Len(t, sharedPool.GetComponentGroups(), 1) require.Same(t, existing, sharedPool.GetComponentGroups()[0]) } func TestGeneratedTiDBGroupCreatesMissingTopologyComponentOnce(t *testing.T) { sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{ CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN, Provider: &apiv1.Provider{Type: apiv1.ProviderType_AWS}, }} tidbGroup := &entity.TiDBGroup{ TiDBGroup: apiv1.TiDBGroup{Id: "tidb-group"}, TopologySpec: &apiv1.TopologySpec{Components: []*apiv1.Component{{ ComponentType: apiv1.ComponentType_TIDB, NodeQuota: &apiv1.NodeQuota{NodeCount: 2, NodeSpec: &apiv1.NodeSpec{Cpu: 4, MemoryGi: 16}}, }}}, } nodeSpecMap := map[string][]*entity.NodeSpec{ "4C16G": {{Cpu: 4, MemoryGi: 16, AvailableCpu: 4, AvailableMemoryMi: 16384}}, } err := (&serverImpl{}).generatedTiDBGroup(context.Background(), sharedPool, &entity.TiDBCluster{}, tidbGroup, nodeSpecMap, nil) require.NoError(t, err) require.Len(t, sharedPool.GetComponentGroups(), 1) require.Equal(t, "tidb-group", sharedPool.GetComponentGroups()[0].GetId()) require.Equal(t, apiv1.Component_CREATING, sharedPool.GetComponentGroups()[0].GetPhase()) } func TestGeneratedTiDBGroupCreatesMissingRCUComponent(t *testing.T) { for _, mode := range []apiv1.CapacityMode{apiv1.CapacityMode_RCU_DRIVEN, apiv1.CapacityMode_CAPACITY_MODE_UNSPECIFIED} { t.Run(mode.String(), func(t *testing.T) { sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{ CapacityMode: mode, Provider: &apiv1.Provider{Type: apiv1.ProviderType_AWS}, }} tidbCluster := &entity.TiDBCluster{RequestUnit: &apiv1.Cluster_RequestUnit{Min: 100, Max: 100}} tidbGroup := &entity.TiDBGroup{TiDBGroup: apiv1.TiDBGroup{Id: "tidb-group"}} nodeSpecMap := map[string][]*entity.NodeSpec{ "4C16G": {{Cpu: 4, MemoryGi: 16, AvailableCpu: 4, AvailableMemoryMi: 16384}}, } ruConfig := &entity.RuConfig{RuConfig: apiv1.RuConfig{Components: map[string]*apiv1.RuConfig_ComponentsMap{ "tidb": {MaxNodeCount: 10, NodeSpecRuRange: map[string]*apiv1.RuConfig_RuRange{ "4c16g": {Min: 100, Max: 0, DefaultNodeCount: uint32Ptr(3)}, }}, }}} err := (&serverImpl{}).generatedTiDBGroup(context.Background(), sharedPool, tidbCluster, tidbGroup, nodeSpecMap, ruConfig) require.NoError(t, err) require.Len(t, sharedPool.GetComponentGroups(), 1) tidbComp := sharedPool.GetComponentGroups()[0] require.Equal(t, apiv1.Component_CREATING, tidbComp.GetPhase()) require.Equal(t, uint32(3), tidbComp.GetNodeQuota().GetNodeCount()) require.NotEmpty(t, tidbComp.GetStorages()) }) } }