package action import ( "context" "testing" "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" "github.com/tidbcloud/cluster-service-ng/config" "github.com/tidbcloud/cluster-service-ng/domain/entity" "github.com/tidbcloud/cluster-service-ng/domain/workflow" "github.com/tidbcloud/cluster-service-ng/pkg/k8s" "github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions" mysqlschema "github.com/tidbcloud/cluster-service-ng/pkg/storage/mysql/schema" repomocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks" templateversion "github.com/tidbcloud/cluster-service-ng/pkg/template/version" "google.golang.org/protobuf/encoding/protojson" apierrors "k8s.io/apimachinery/pkg/api/errors" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/types" "sigs.k8s.io/controller-runtime/pkg/client/fake" ) func TestShouldSkipFullComponentConfigForDeletingComponent(t *testing.T) { tests := []struct { name string comp *entity.Component want bool }{ { name: "deleting tikv group skips for delete replay tolerance", comp: &entity.Component{ ComponentType: apiv1.ComponentType_TIKV, Phase: apiv1.Component_DELETING, }, want: true, }, { name: "deleting coprocessor worker still hydrates config", comp: &entity.Component{ ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER, Phase: apiv1.Component_DELETING, }, }, { name: "available tikv still hydrates config", comp: &entity.Component{ ComponentType: apiv1.ComponentType_TIKV, Phase: apiv1.Component_MODIFYING, }, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { assert.Equal(t, tt.want, shouldSkipFullComponentConfigForDeletingComponent(tt.comp)) }) } } func TestApplyComponentGroupDeletesTiKVGroupDirectlyWithExpiredPinnedTemplateVersion(t *testing.T) { ctx := context.Background() scheme := runtime.NewScheme() require.NoError(t, v1alpha1.AddToScheme(scheme)) comp := &entity.Component{ Id: "blue", ComponentType: apiv1.ComponentType_TIKV, Phase: apiv1.Component_DELETING, TemplateVersion: "expired-version", } sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "sp-1", DataplaneName: "dp-1", ComponentGroups: []*apiv1.Component{comp}, }, } sharedPoolBytes, err := protojson.Marshal(sharedPool) require.NoError(t, err) tikvGroup := &v1alpha1.TiKVGroup{ ObjectMeta: metav1.ObjectMeta{ Name: entity.GetComponentGroupName(comp), Namespace: sharedPool.GetDataPlaneNamespace(), Annotations: map[string]string{ templateversion.TemplateVersionAnnotationKey: "expired-version", }, }, } kubeClient := fake.NewClientBuilder().WithScheme(scheme).WithObjects(tikvGroup).Build() sharedPoolRepo := repomocks.NewSharedPoolRepo(t) sharedPoolRepo.On("Load", mock.Anything, sharedPool.GetId(), false).Return(sharedPool, nil).Once() templateRepo := repomocks.NewTemplateRegistryRepo(t) templateRepo.On("Load", mock.Anything, "tikv_group", "expired-version").Return(&entity.TemplateRegistry{ TemplateVersionInfo: entity.TemplateVersionInfo{ TemplateName: "tikv_group", Version: "expired-version", Status: apiv1.TemplateVersionStatus_TEMPLATE_VERSION_STATUS_EXPIRED, }, Content: "kind: TiKVGroup\n", }, nil).Maybe() action := &ApplyComponentGroupAction{ sharedPoolRepo: sharedPoolRepo, dataplaneManager: &mockDataplaneManager{client: kubeClient}, templateVersionSvc: templateversion.NewService(templateRepo, 3), } err = action.Invoke(actions.ActionArgs{ workflow.SharedPoolIdKey: actions.ActionArgString(sharedPool.GetId()), workflow.SharedPoolKey: actions.ActionArgString(string(sharedPoolBytes)), workflow.ComponentIdsKey: actions.ActionArgString(comp.GetId()), workflow.ComponentTypesKey: actions.ActionArgString(entity.GetComponentTypeName(comp)), }) require.NoError(t, err) stored := &v1alpha1.TiKVGroup{} err = kubeClient.Get(ctx, types.NamespacedName{Name: tikvGroup.Name, Namespace: tikvGroup.Namespace}, stored) assert.True(t, apierrors.IsNotFound(err)) } func TestBuildCompTempArgs_ServerLabels(t *testing.T) { imageConfig := &config.ImageConfig{ Registry: "registry.example.com", RepoPrefix: "tidb", ImagePullPolicy: "IfNotPresent", } sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "shared-pool-1", Provider: &apiv1.Provider{ NodeZones: []string{"zone-1", "zone-2"}, }, }, } nodeVersionMap := map[string]string{ "4C16G": "v1", } tests := []struct { name string component *entity.Component expectLabelsKey bool expectLabelsValue string }{ { name: "component without labels", component: &entity.Component{ Id: "comp-1", ComponentType: apiv1.ComponentType_TIKV, Version: "v7.5.0", Labels: nil, Configuration: &apiv1.Component_Configuration{ ImageRepo: "pingcap/tikv", ImageTag: "v7.5.0", }, NodeQuota: &apiv1.NodeQuota{ NodeCount: 3, MaxNodeCount: 5, MinNodeCount: 1, NodeSpec: &apiv1.NodeSpec{ Cpu: 4, MemoryGi: 16, AvailableCpu: 4, AvailableMemoryMi: 16384, }, }, }, expectLabelsKey: false, expectLabelsValue: "", }, { name: "component with empty labels", component: &entity.Component{ Id: "comp-2", ComponentType: apiv1.ComponentType_TIKV, Version: "v7.5.0", Labels: map[string]string{}, Configuration: &apiv1.Component_Configuration{ ImageRepo: "pingcap/tikv", ImageTag: "v7.5.0", }, NodeQuota: &apiv1.NodeQuota{ NodeCount: 3, MaxNodeCount: 5, MinNodeCount: 1, NodeSpec: &apiv1.NodeSpec{ Cpu: 4, MemoryGi: 16, AvailableCpu: 4, AvailableMemoryMi: 16384, }, }, }, expectLabelsKey: false, expectLabelsValue: "", }, { name: "component with single label", component: &entity.Component{ Id: "comp-3", ComponentType: apiv1.ComponentType_TIKV, Version: "v7.5.0", Labels: map[string]string{ "keyspace": "cluster-123", }, Configuration: &apiv1.Component_Configuration{ ImageRepo: "pingcap/tikv", ImageTag: "v7.5.0", }, NodeQuota: &apiv1.NodeQuota{ NodeCount: 3, MaxNodeCount: 5, MinNodeCount: 1, NodeSpec: &apiv1.NodeSpec{ Cpu: 4, MemoryGi: 16, AvailableCpu: 4, AvailableMemoryMi: 16384, }, }, }, expectLabelsKey: true, expectLabelsValue: "labels:\n \"keyspace\": \"cluster-123\"", }, { name: "component with multiple labels sorted", component: &entity.Component{ Id: "comp-4", ComponentType: apiv1.ComponentType_TIKV, Version: "v7.5.0", Labels: map[string]string{ "keyspace": "cluster-123", "cidr": "10.0.0.0/8", "custom": "value", }, Configuration: &apiv1.Component_Configuration{ ImageRepo: "pingcap/tikv", ImageTag: "v7.5.0", }, NodeQuota: &apiv1.NodeQuota{ NodeCount: 3, MaxNodeCount: 5, MinNodeCount: 1, NodeSpec: &apiv1.NodeSpec{ Cpu: 4, MemoryGi: 16, AvailableCpu: 4, AvailableMemoryMi: 16384, }, }, }, expectLabelsKey: true, expectLabelsValue: "labels:\n \"cidr\": \"10.0.0.0/8\"\n \"custom\": \"value\"\n \"keyspace\": \"cluster-123\"", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { templArgs, err := buildCompTempArgs(imageConfig, tt.component, sharedPool, nodeVersionMap) require.NoError(t, err) labelsValue, hasLabels := templArgs["ServerLabels"] if tt.expectLabelsKey { assert.True(t, hasLabels, "ServerLabels key should exist in templArgs") assert.Equal(t, tt.expectLabelsValue, labelsValue, "ServerLabels value mismatch") } else { assert.False(t, hasLabels, "ServerLabels key should not exist in templArgs when labels are empty or nil") } }) } } func TestBuildCompTempArgs_BasicFields(t *testing.T) { imageConfig := &config.ImageConfig{ Registry: "registry.example.com", RepoPrefix: "tidb", ImagePullPolicy: "Always", } sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "shared-pool-1", Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, }, } nodeVersionMap := map[string]string{ "4C16G": "v2", } component := &entity.Component{ Id: "comp-test", ComponentType: apiv1.ComponentType_TIKV, Version: "v7.5.0", Configuration: &apiv1.Component_Configuration{ ImageRepo: "pingcap/tikv", ImageTag: "v7.5.0", ConfigurationFileContent: "[tikv]\nlog-level = \"info\"", }, NodeQuota: &apiv1.NodeQuota{ NodeCount: 3, MaxNodeCount: 5, MinNodeCount: 1, NodeSpec: &apiv1.NodeSpec{ Cpu: 4, MemoryGi: 16, AvailableCpu: 4, AvailableMemoryMi: 16384, }, }, } templArgs, err := buildCompTempArgs(imageConfig, component, sharedPool, nodeVersionMap) require.NoError(t, err) // Verify basic fields assert.Equal(t, "registry.example.com/tidb/tikv:v7.5.0", templArgs["Image"]) assert.Equal(t, "shared-pool-1", templArgs["ClusterId"]) assert.Equal(t, "shared-pool-1", templArgs["SharedPoolId"]) assert.Equal(t, "comp-test", templArgs["ComponentId"]) assert.Equal(t, "v7.5.0", templArgs["Version"]) assert.Equal(t, "3", templArgs["Replicas"]) assert.Equal(t, "Always", templArgs["ImagePullPolicy"]) assert.Equal(t, "5", templArgs["MaxNodeCount"]) assert.Equal(t, "1", templArgs["MinNodeCount"]) assert.Equal(t, "v2", templArgs["NodeVersion"]) assert.Equal(t, "[tikv]\nlog-level = \"info\"", templArgs["Config"]) assert.NotContains(t, templArgs, "TiProxyGracefulShutdownDeleteDelaySeconds") assert.NotContains(t, templArgs, "TiProxyGracefulRestartMaxSurge") } func TestSetTiKVPlacementExclusiveTemplateArg(t *testing.T) { tests := []struct { name string args actions.ActionArgs wantSet bool wantErrText string }{ { name: "exclusive true", args: actions.ActionArgs{workflow.TiKVPlacementExclusiveKey: actions.ActionArgBool(true)}, wantSet: true, }, { name: "exclusive false", args: actions.ActionArgs{workflow.TiKVPlacementExclusiveKey: actions.ActionArgBool(false)}, wantSet: false, }, { name: "missing arg", args: actions.ActionArgs{}, wantSet: false, }, { name: "invalid arg", args: actions.ActionArgs{workflow.TiKVPlacementExclusiveKey: actions.ActionArgString("not-bool")}, wantErrText: "failed to parse tikv_placement_exclusive", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { templArgs := map[string]string{} err := setTiKVPlacementExclusiveTemplateArg(tt.args, templArgs) if tt.wantErrText != "" { require.Error(t, err) assert.Contains(t, err.Error(), tt.wantErrText) return } require.NoError(t, err) got, ok := templArgs["TiKVPlacementExclusive"] if tt.wantSet { require.True(t, ok) assert.Equal(t, "true", got) } else { assert.False(t, ok) } }) } } func TestRenderTiKVGroupPlacementExclusive(t *testing.T) { tl := initTemplateLibrary(t, []string{"tikv_group"}) baseArgs := map[string]string{ "ClusterId": "123", "SharedPoolId": "sp-1", "ComponentType": "tikv", "ComponentGroup": "db-tikv", "NodeSpec": "4C16G", "MinNodeCount": "1", "MaxNodeCount": "3", "Replicas": "3", "Topologies": "- topologyKey: topology.kubernetes.io/zone", "Version": "v9.0.0", "Image": "tikv:v9.0.0", "Config": "[tikv]", "DataStorageSize": "100", "DataStorageClass": "gp3", "TiKVPreStopImage": "prestop:v1", "DfsPrefix": "prefix", "DfsBucket": "bucket", "DfsRegion": "us-west-2", "ImagePullPolicy": "IfNotPresent", "InitImage": "busybox:1.0", "NodeVersion": "v1", } yamlStr, err := tl.Render("tikv_group", baseArgs) require.NoError(t, err) obj, _, err := k8s.GenerateUnstructuredFromYaml(yamlStr) require.NoError(t, err) _, found, err := unstructured.NestedMap(obj.Object, "spec", "template", "spec", "placement") require.NoError(t, err) assert.False(t, found) exclusiveArgs := copyStringMap(baseArgs) exclusiveArgs["TiKVPlacementExclusive"] = "true" yamlStr, err = tl.Render("tikv_group", exclusiveArgs) require.NoError(t, err) obj, _, err = k8s.GenerateUnstructuredFromYaml(yamlStr) require.NoError(t, err) exclusive, found, err := unstructured.NestedBool(obj.Object, "spec", "template", "spec", "placement", "exclusive") require.NoError(t, err) require.True(t, found) assert.True(t, exclusive) } func TestSetTiFlashColumnarTemplateArg(t *testing.T) { tests := []struct { name string tiflashMode string wantSet bool }{ {name: "marker pool", tiflashMode: mysqlschema.SharedPoolTiflashModeColumnar, wantSet: true}, {name: "legacy pool", wantSet: false}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { templArgs := map[string]string{} sharedPool := &entity.SharedPool{TiflashMode: tt.tiflashMode} setTiFlashColumnarTemplateArg(templArgs, sharedPool) if tt.wantSet { assert.Equal(t, "true", templArgs["TiFlashColumnar"]) } else { assert.NotContains(t, templArgs, "TiFlashColumnar") } }) } } func TestBuildCompTempArgs_TiProxyGracefulRestartArgs(t *testing.T) { imageConfig := &config.ImageConfig{ Registry: "registry.example.com", RepoPrefix: "tidb", ImagePullPolicy: "IfNotPresent", } nodeVersionMap := map[string]string{ "4C16G": "v1", } tests := []struct { name string componentType apiv1.ComponentType wantSet bool }{ { name: "enabled by default for tiproxy", componentType: apiv1.ComponentType_TIPROXY, wantSet: true, }, { name: "not set for non-tiproxy component", componentType: apiv1.ComponentType_TIKV, wantSet: false, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "shared-pool-1", Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, }, } component := &entity.Component{ Id: "comp-test", ComponentType: tt.componentType, Version: "v7.5.0", Configuration: &apiv1.Component_Configuration{ ImageRepo: "pingcap/tiproxy", ImageTag: "v7.5.0", }, NodeQuota: &apiv1.NodeQuota{ NodeCount: 3, MaxNodeCount: 5, MinNodeCount: 1, NodeSpec: &apiv1.NodeSpec{ Cpu: 4, MemoryGi: 16, AvailableCpu: 4, AvailableMemoryMi: 16384, }, }, } templArgs, err := buildCompTempArgs(imageConfig, component, sharedPool, nodeVersionMap) require.NoError(t, err) if tt.wantSet { assert.Equal(t, "86400", templArgs["TiProxyGracefulShutdownDeleteDelaySeconds"]) assert.Equal(t, "1024", templArgs["TiProxyGracefulRestartMaxSurge"]) } else { assert.NotContains(t, templArgs, "TiProxyGracefulShutdownDeleteDelaySeconds") assert.NotContains(t, templArgs, "TiProxyGracefulRestartMaxSurge") } }) } } func TestBuildCompTempArgs_DefaultImagePullPolicy(t *testing.T) { imageConfig := &config.ImageConfig{ Registry: "registry.example.com", RepoPrefix: "tidb", ImagePullPolicy: "", // empty, should default to IfNotPresent } sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "shared-pool-1", Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, }, } nodeVersionMap := map[string]string{ "4C16G": "v1", } component := &entity.Component{ Id: "comp-test", ComponentType: apiv1.ComponentType_TIKV, Version: "v7.5.0", Configuration: &apiv1.Component_Configuration{ ImageRepo: "pingcap/tikv", ImageTag: "v7.5.0", }, NodeQuota: &apiv1.NodeQuota{ NodeCount: 3, NodeSpec: &apiv1.NodeSpec{ Cpu: 4, MemoryGi: 16, AvailableCpu: 4, AvailableMemoryMi: 16384, }, }, } templArgs, err := buildCompTempArgs(imageConfig, component, sharedPool, nodeVersionMap) require.NoError(t, err) assert.Equal(t, "IfNotPresent", templArgs["ImagePullPolicy"]) } func TestBuildCompTempArgs_DeleteSkipsAvailableResourceValidation(t *testing.T) { imageConfig := &config.ImageConfig{ Registry: "registry.example.com", RepoPrefix: "tidb", ImagePullPolicy: "IfNotPresent", } sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "shared-pool-1", Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, }, } component := &entity.Component{ Id: "comp-delete", ComponentType: apiv1.ComponentType_TIDB, Version: "v8.5.4", Phase: apiv1.Component_DELETING, Configuration: &apiv1.Component_Configuration{ ImageRepo: "pingcap/tidb", ImageTag: "v8.5.4", }, NodeQuota: &apiv1.NodeQuota{ NodeCount: 0, MaxNodeCount: 1, MinNodeCount: 1, NodeSpec: &apiv1.NodeSpec{ Cpu: 4, MemoryGi: 8, }, }, } templArgs, err := buildCompTempArgs(imageConfig, component, sharedPool, map[string]string{}) require.NoError(t, err) assert.Equal(t, "v1.0.0", templArgs["NodeVersion"]) assert.Equal(t, "0", templArgs["Replicas"]) assert.Equal(t, "6553", templArgs["GOMEMLIMIT"]) _, hasSetResource := templArgs["SetResource"] assert.False(t, hasSetResource, "SetResource should not be set for deleting components") _, hasMemory := templArgs["Memory"] assert.False(t, hasMemory, "Memory should not be required for deleting components") tl := initTemplateLibrary(t, []string{"tidb_group"}) templArgs["UserClusterId"] = "cluster-1" templArgs["UserTenantId"] = "tenant-1" templArgs["SlowLogImage"] = "busybox:1.0" templArgs["AuditLogImage"] = "audit-ng:1.0" templArgs["SlowLogStoragePath"] = "/var/log/tidb" templArgs["SlowLogStorageSize"] = "10" templArgs["SlowLogStorageClass"] = "standard" templArgs["SlowLogVolumeAttributesClass"] = "vac" templArgs["PreStopSleepSeconds"] = "35" templArgs["AuditLogConfig"] = "logDir: /var/log/auditlog" templArgs["TiDBKeyspace"] = "cluster-1" yamlStr, err := tl.Render("tidb_group", templArgs) require.NoError(t, err) assert.NotContains(t, yamlStr, "{{!GOMEMLIMIT}}") _, _, err = k8s.GenerateUnstructuredFromYaml(yamlStr) require.NoError(t, err) } func TestBuildCompTempArgs_DeleteStubWithoutNodeSpecRendersTiDBGroup(t *testing.T) { imageConfig := &config.ImageConfig{ Registry: "registry.example.com", RepoPrefix: "tidb", ImagePullPolicy: "IfNotPresent", } sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "shared-pool-1", Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, }, } component := &entity.Component{ Id: "comp-delete", ComponentType: apiv1.ComponentType_TIDB, Version: "v8.5.4", Phase: apiv1.Component_DELETING, Configuration: &apiv1.Component_Configuration{ ImageRepo: "pingcap/tidb", ImageTag: "v8.5.4", ConfigurationFileContent: "", }, NodeQuota: &apiv1.NodeQuota{ NodeCount: 0, }, } templArgs, err := buildCompTempArgs(imageConfig, component, sharedPool, map[string]string{}) require.NoError(t, err) assert.Equal(t, "819", templArgs["GOMEMLIMIT"]) tl := initTemplateLibrary(t, []string{"tidb_group"}) templArgs["UserClusterId"] = "cluster-1" templArgs["UserTenantId"] = "tenant-1" templArgs["SlowLogImage"] = "busybox:1.0" templArgs["AuditLogImage"] = "audit-ng:1.0" templArgs["SlowLogStoragePath"] = entity.AuditAndSlowLogPath templArgs["SlowLogStorageSize"] = "1" templArgs["SlowLogStorageClass"] = "delete-placeholder" templArgs["SlowLogVolumeAttributesClass"] = "delete-placeholder" templArgs["PreStopSleepSeconds"] = "30" templArgs["AuditLogConfig"] = "logDir: /var/log/auditlog" templArgs["TiDBKeyspace"] = "cluster-1" yamlStr, err := tl.Render("tidb_group", templArgs) require.NoError(t, err) assert.NotContains(t, yamlStr, "{{!GOMEMLIMIT}}") _, _, err = k8s.GenerateUnstructuredFromYaml(yamlStr) require.NoError(t, err) } func TestInjectPDLocationLabels(t *testing.T) { // PG eligibility is read from the shared pool's durable per-pool gate // (Provider.PlacementGroupEnabled, set at CreateSharedPool time) so the // scheduling side and the NodePool infrastructure side never diverge on // the same pool. A shared pool with gate=true keeps this test on the // "PG enabled" path; a legacy pool with gate=false (even on single AZ) // skips it. sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Provider: &apiv1.Provider{NodeZones: []string{"us-west-2a"}, PlacementGroupEnabled: true}, }, } tests := []struct { name string compType apiv1.ComponentType config string expect string }{ { name: "pd appends replication section", compType: apiv1.ComponentType_PD, config: "", expect: "\n[replication]\n location-labels = [\"tidbcloud.com/pg-partition\"]", }, { name: "pd scheduling inserts existing replication section", compType: apiv1.ComponentType_PD_SCHEDULING, config: "[replication]\nmax-replicas = 3", expect: "[replication]\n location-labels = [\"tidbcloud.com/pg-partition\"]\nmax-replicas = 3", }, { name: "existing location labels unchanged", compType: apiv1.ComponentType_PD, config: "[replication]\n location-labels = [\"zone\"]", expect: "[replication]\n location-labels = [\"zone\"]", }, { name: "pd tso unchanged", compType: apiv1.ComponentType_PD_TSO, config: "", expect: "", }, { name: "resource manager unchanged", compType: apiv1.ComponentType_RESOURCE_MANAGER, config: "", expect: "", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { templArgs := map[string]string{"Config": tt.config} comp := &entity.Component{ComponentType: tt.compType} injectPDLocationLabels(comp, sharedPool, templArgs) assert.Equal(t, tt.expect, templArgs["Config"]) }) } } func TestInjectPDLocationLabelsDisabled(t *testing.T) { sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Provider: &apiv1.Provider{}}} templArgs := map[string]string{"Config": ""} comp := &entity.Component{ComponentType: apiv1.ComponentType_PD} injectPDLocationLabels(comp, sharedPool, templArgs) assert.Empty(t, templArgs["Config"]) } // TestInjectPDLocationLabelsLegacyPGDisabledKeepsGateClosed is the regression // test for the silent-mode-flip bug: a legacy single-AZ shared pool with // Provider.PlacementGroupEnabled=false (set by an older build) must NOT // receive pg-partition scheduling, even though its node-zone count would // otherwise read as single-AZ. The gate is the per-pool migration marker, // not the live zone list. func TestInjectPDLocationLabelsLegacyPGDisabledKeepsGateClosed(t *testing.T) { sharedPool := &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Provider: &apiv1.Provider{ NodeZones: []string{"us-west-2a"}, PlacementGroupEnabled: false, }, }, } templArgs := map[string]string{"Config": ""} comp := &entity.Component{ComponentType: apiv1.ComponentType_PD} injectPDLocationLabels(comp, sharedPool, templArgs) assert.Empty(t, templArgs["Config"], "legacy PG-disabled shared pool must not get pg-partition scheduling") } func TestBuildCompTempArgs_PDMode(t *testing.T) { imageConfig := &config.ImageConfig{ Registry: "registry.example.com", RepoPrefix: "tidb", ImagePullPolicy: "IfNotPresent", } nodeVersionMap := map[string]string{ "4C16G": "v1", } basePD := &entity.Component{ Id: "pd-test", ComponentType: apiv1.ComponentType_PD, Version: "v8.5.0", Configuration: &apiv1.Component_Configuration{ ImageRepo: "pingcap/pd", ImageTag: "v8.5.0", }, NodeQuota: &apiv1.NodeQuota{ NodeCount: 3, NodeSpec: &apiv1.NodeSpec{ Cpu: 4, MemoryGi: 16, AvailableCpu: 4, AvailableMemoryMi: 16384, }, }, } tests := []struct { name string sharedPool *entity.SharedPool expectMode string expectModeField bool }{ { name: "topology driven with pd_tso node_count=0 uses normal pd mode", sharedPool: &entity.SharedPool{SharedPool: apiv1.SharedPool{ Id: "sp-topo-no-tso", CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN, Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, ComponentGroups: []*apiv1.Component{ {ComponentType: apiv1.ComponentType_PD}, {ComponentType: apiv1.ComponentType_TIKV}, {ComponentType: apiv1.ComponentType_PD_TSO, NodeQuota: &apiv1.NodeQuota{NodeCount: 0}}, }, }}, expectMode: "", expectModeField: false, }, { name: "topology driven with pd_tso uses ms pd mode", sharedPool: &entity.SharedPool{SharedPool: apiv1.SharedPool{ Id: "sp-topo-with-tso", CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN, Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, ComponentGroups: []*apiv1.Component{ {ComponentType: apiv1.ComponentType_PD}, {ComponentType: apiv1.ComponentType_PD_TSO, NodeQuota: &apiv1.NodeQuota{NodeCount: 2}}, }, }}, expectMode: "ms", expectModeField: true, }, { name: "rcu driven keeps ms pd mode", sharedPool: &entity.SharedPool{SharedPool: apiv1.SharedPool{ Id: "sp-rcu", CapacityMode: apiv1.CapacityMode_RCU_DRIVEN, Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, }}, expectMode: "ms", expectModeField: true, }, { name: "monolithic rcu driven omits pd mode", sharedPool: &entity.SharedPool{SharedPool: apiv1.SharedPool{ Id: "sp-mono", CapacityMode: apiv1.CapacityMode_RCU_DRIVEN, PdDeploymentMode: apiv1.PDDeploymentMode_PD_DEPLOYMENT_MODE_MONOLITHIC, Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, }}, expectMode: "", expectModeField: false, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { templArgs, err := buildCompTempArgs(imageConfig, basePD, tt.sharedPool, nodeVersionMap) require.NoError(t, err) if tt.expectModeField { assert.Equal(t, tt.expectMode, templArgs["PDMode"]) return } _, ok := templArgs["PDMode"] assert.False(t, ok) }) } } func TestResolvePDMode_ZeroNodeCountTreatAsDisabled(t *testing.T) { sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{ CapacityMode: apiv1.CapacityMode_TOPOLOGY_DRIVEN, ComponentGroups: []*apiv1.Component{ {ComponentType: apiv1.ComponentType_PD}, {ComponentType: apiv1.ComponentType_PD_TSO, NodeQuota: &apiv1.NodeQuota{NodeCount: 0}}, }, }} mode, ok := resolvePDMode(sharedPool) assert.False(t, ok) assert.Empty(t, mode) } func TestDesiredSpecForUpdate_PreservesDeploymentReplicasOnlyWhenTemplateOmitsIt(t *testing.T) { existing := &unstructured.Unstructured{ Object: map[string]any{ "kind": "Deployment", "spec": map[string]any{ "replicas": int64(5), }, }, } t.Run("non deployment returns desired spec", func(t *testing.T) { desired := &unstructured.Unstructured{ Object: map[string]any{ "kind": "StatefulSet", "spec": map[string]any{ "template": "desired", }, }, } assert.Equal(t, desired.Object["spec"], desiredSpecForUpdate(existing, desired)) }) t.Run("deployment with desired replicas keeps desired value", func(t *testing.T) { desired := &unstructured.Unstructured{ Object: map[string]any{ "kind": "Deployment", "spec": map[string]any{ "replicas": int64(2), "template": "desired", }, }, } spec := desiredSpecForUpdate(existing, desired) assert.Equal(t, desired.Object["spec"], spec) assert.Equal(t, int64(2), spec.(map[string]any)["replicas"]) }) t.Run("deployment without desired replicas preserves live value", func(t *testing.T) { desired := &unstructured.Unstructured{ Object: map[string]any{ "kind": "Deployment", "spec": map[string]any{ "template": "desired", }, }, } spec := desiredSpecForUpdate(existing, desired).(map[string]any) assert.Equal(t, "desired", spec["template"]) assert.Equal(t, int64(5), spec["replicas"]) }) t.Run("tikv group without desired placement preserves live value", func(t *testing.T) { current := &unstructured.Unstructured{ Object: map[string]any{ "kind": "TiKVGroup", "spec": map[string]any{ "template": map[string]any{ "spec": map[string]any{ "placement": map[string]any{ "exclusive": true, }, }, }, }, }, } desired := &unstructured.Unstructured{ Object: map[string]any{ "kind": "TiKVGroup", "spec": map[string]any{ "template": map[string]any{ "spec": map[string]any{ "version": "v9.0.0", }, }, }, }, } spec := desiredSpecForUpdate(current, desired).(map[string]any) exclusive, found, err := unstructured.NestedBool(spec, "template", "spec", "placement", "exclusive") require.NoError(t, err) require.True(t, found) assert.True(t, exclusive) }) t.Run("tikv group with desired placement keeps desired value", func(t *testing.T) { current := &unstructured.Unstructured{ Object: map[string]any{ "kind": "TiKVGroup", "spec": map[string]any{ "template": map[string]any{ "spec": map[string]any{ "placement": map[string]any{ "exclusive": true, }, }, }, }, }, } desired := &unstructured.Unstructured{ Object: map[string]any{ "kind": "TiKVGroup", "spec": map[string]any{ "template": map[string]any{ "spec": map[string]any{ "placement": map[string]any{ "exclusive": false, }, }, }, }, }, } spec := desiredSpecForUpdate(current, desired).(map[string]any) exclusive, found, err := unstructured.NestedBool(spec, "template", "spec", "placement", "exclusive") require.NoError(t, err) require.True(t, found) assert.False(t, exclusive) }) } func TestSetTiKVRemoteWorkerTemplateArgs(t *testing.T) { tests := []struct { name string sharedPool *entity.SharedPool wantRemoteWorker string wantRemoteCoprocessor string wantNoRemoteWorkers bool wantErr bool }{ { name: "legacy empty set with dedicated positive points to dedicated", sharedPool: &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-1", Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, ComponentGroups: []*apiv1.Component{ {Id: "pool-1", ComponentType: apiv1.ComponentType_TIKV_WORKER, NodeQuota: &apiv1.NodeQuota{NodeCount: 1}}, {Id: "pool-1", ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER, NodeQuota: &apiv1.NodeQuota{NodeCount: 1}}, }, }, }, wantRemoteWorker: "db-pool-1-default", wantRemoteCoprocessor: "db-pool-1-coprocessor", }, { name: "disabled falls back to default", sharedPool: &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-2", EnabledTikvRemoteWorkers: []apiv1.ComponentType{apiv1.ComponentType_TIKV_WORKER}, Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, ComponentGroups: []*apiv1.Component{ {Id: "pool-2", ComponentType: apiv1.ComponentType_TIKV_WORKER, NodeQuota: &apiv1.NodeQuota{NodeCount: 1}}, {Id: "pool-2", ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER, NodeQuota: &apiv1.NodeQuota{NodeCount: 1}}, }, }, }, wantRemoteWorker: "db-pool-2-default", wantRemoteCoprocessor: "db-pool-2-default", }, { name: "zero dedicated falls back to default", sharedPool: &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-3", Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, ComponentGroups: []*apiv1.Component{ {Id: "pool-3", ComponentType: apiv1.ComponentType_TIKV_WORKER, NodeQuota: &apiv1.NodeQuota{NodeCount: 1}}, {Id: "pool-3", ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER, NodeQuota: &apiv1.NodeQuota{NodeCount: 0}}, }, }, }, wantRemoteWorker: "db-pool-3-default", wantRemoteCoprocessor: "db-pool-3-default", }, { name: "missing dedicated falls back to default", sharedPool: &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-4", EnabledTikvRemoteWorkers: []apiv1.ComponentType{apiv1.ComponentType_TIKV_WORKER, apiv1.ComponentType_COPROCESSOR_WORKER}, Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, ComponentGroups: []*apiv1.Component{ {Id: "pool-4", ComponentType: apiv1.ComponentType_TIKV_WORKER, NodeQuota: &apiv1.NodeQuota{NodeCount: 1}}, }, }, }, wantRemoteWorker: "db-pool-4-default", wantRemoteCoprocessor: "db-pool-4-default", }, { name: "missing default does not render remoteWorkers", sharedPool: &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-5", EnabledTikvRemoteWorkers: []apiv1.ComponentType{apiv1.ComponentType_TIKV_WORKER}, Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, ComponentGroups: []*apiv1.Component{ {Id: "pool-5", ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER, NodeQuota: &apiv1.NodeQuota{NodeCount: 1}}, }, }, }, wantNoRemoteWorkers: true, }, { name: "zero default does not render remoteWorkers", sharedPool: &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-6", EnabledTikvRemoteWorkers: []apiv1.ComponentType{apiv1.ComponentType_TIKV_WORKER}, Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, ComponentGroups: []*apiv1.Component{ {Id: "pool-6", ComponentType: apiv1.ComponentType_TIKV_WORKER, NodeQuota: &apiv1.NodeQuota{NodeCount: 0}}, {Id: "pool-6", ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER, NodeQuota: &apiv1.NodeQuota{NodeCount: 1}}, }, }, }, wantNoRemoteWorkers: true, }, { name: "invalid enabled set returns error and does not write args", sharedPool: &entity.SharedPool{ SharedPool: apiv1.SharedPool{ Id: "pool-7", EnabledTikvRemoteWorkers: []apiv1.ComponentType{apiv1.ComponentType_COPROCESSOR_WORKER}, Provider: &apiv1.Provider{ NodeZones: []string{"zone-1"}, }, ComponentGroups: []*apiv1.Component{ {Id: "pool-7", ComponentType: apiv1.ComponentType_TIKV_WORKER, NodeQuota: &apiv1.NodeQuota{NodeCount: 1}}, {Id: "pool-7", ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER, NodeQuota: &apiv1.NodeQuota{NodeCount: 1}}, }, }, }, wantErr: true, wantNoRemoteWorkers: true, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { templArgs := make(map[string]string) err := setTiKVRemoteWorkerTemplateArgs(templArgs, tt.sharedPool) if tt.wantErr { require.Error(t, err) } else { require.NoError(t, err) } if tt.wantNoRemoteWorkers { assert.NotContains(t, templArgs, "RemoteWorkerGroup") assert.NotContains(t, templArgs, "RemoteCoprocessorWorkerGroup") } else { assert.Equal(t, tt.wantRemoteWorker, templArgs["RemoteWorkerGroup"]) assert.Equal(t, tt.wantRemoteCoprocessor, templArgs["RemoteCoprocessorWorkerGroup"]) } }) } } func TestBuildCompTempArgs_DeletingCoprocessorRepairSnapshot(t *testing.T) { imageConfig := &config.ImageConfig{Registry: "registry.example.com", RepoPrefix: "tidb", ImagePullPolicy: "IfNotPresent"} sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{Id: "shared-pool-1", Provider: &apiv1.Provider{NodeZones: []string{"zone-1"}}}} component := &entity.Component{ Id: "default", ComponentType: apiv1.ComponentType_COPROCESSOR_WORKER, Phase: apiv1.Component_DELETING, NodeQuota: &apiv1.NodeQuota{}, } templArgs, err := buildCompTempArgs(imageConfig, component, sharedPool, map[string]string{}) require.NoError(t, err) assert.Equal(t, "0", templArgs["Replicas"]) assert.Equal(t, "v1.0.0", templArgs["NodeVersion"]) }