package server

import (
	"context"
	"fmt"
	"testing"

	. "github.com/onsi/gomega"
	"github.com/samber/lo"
	"github.com/stretchr/testify/assert"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/yaml"

	"github.com/tidbcloud/cluster-controller/server/internal/util"
	"github.com/tidbcloud/cluster-controller/server/pkg/cloudprovider"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-server/v1"
	apicommon "github.com/tidbcloud/infra-provider/api/tidbcloud/common"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/gitops"
	"github.com/tidbcloud/infra-provider/pkg/scheme"
)

const (
	existClusterId        = "123456"
	defaultLoadBalancerID = "666"
)

func buildInstanceTypeConfigMap() *corev1.ConfigMap {
	instanceTypes := []cloudv1.InstanceType{
		{
			Name:   "",
			CPU:    resource.MustParse("0"),
			Memory: resource.MustParse("0"),
		},
		{
			Name:   "2c4g",
			CPU:    resource.MustParse("2"),
			Memory: resource.MustParse("4Gi"),
		},
		{
			Name:   "4c8g",
			CPU:    resource.MustParse("4"),
			Memory: resource.MustParse("8Gi"),
		},
		{
			Name:   "8c16g",
			CPU:    resource.MustParse("8"),
			Memory: resource.MustParse("16Gi"),
		},
	}

	yamlData, _ := yaml.Marshal(instanceTypes)
	return &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "instance-type-config",
			Namespace: "cluster-service",
		},
		Data: map[string]string{
			"instancetype-config.yaml": string(yamlData),
		},
	}
}

func buildTestClusterCR(clusterId string, apply func(*cloudv1.Cluster)) *cloudv1.Cluster {
	cr := &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name: clusterId,
			Labels: map[string]string{
				cloudv1.ClusterTagKey: clusterId,
			},
			Annotations: map[string]string{
				consts.DefaultLoadBalancerIDAnnotationKey: defaultLoadBalancerID,
			},
		},
		Spec: cloudv1.ClusterSpec{
			ConfigRevisionRef: &cloudv1.Reference{
				Name: "rev1",
			},
		},
		Status: cloudv1.ClusterStatus{
			Phase: cloudv1.ClusterPhaseRunning,
			DB:    &cloudv1.DBStatus{},
		},
	}

	if apply != nil {
		apply(cr)
	}

	// To make `cluster.Synced()` return `true`
	if cr.Spec.Cluster.TiDB.Replicas != 0 {
		members := make([]cloudv1.TiDBMember, cr.Spec.Cluster.TiDB.Replicas)
		if cr.Spec.Cluster.TiDB.InstanceType != "" {
			for i := range members {
				members[i].InstanceType = cr.Spec.Cluster.TiDB.InstanceType
			}
		}
		cr.Status.DB.TiDB = &cloudv1.TiDBStatus{
			Status:  cloudv1.ComponentNormal,
			Members: members,
		}
	}

	if cr.Spec.Cluster.TiKV.Replicas != 0 {
		members := make([]cloudv1.StoreMember, cr.Spec.Cluster.TiKV.Replicas)
		if cr.Spec.Cluster.TiKV.InstanceType != "" {
			for i := range members {
				members[i].InstanceType = cr.Spec.Cluster.TiKV.InstanceType
			}
		}
		storageVolumes := cr.Status.DB.TiKV.StorageVolumes
		cr.Status.DB.TiKV = &cloudv1.TiKVStatus{
			Status:         cloudv1.ComponentNormal,
			Members:        members,
			StorageVolumes: storageVolumes,
		}
	}

	// Make all extra groups synced
	cr.Status.ExtraEndpoints = make([]cloudv1.ExtraGroupStatus, 0, len(cr.Spec.ExtraGroups))
	for _, extraGroup := range cr.GetExtraGroups() {
		if extraGroup.Spec.TiDB.Replicas != 0 {
			members := make([]cloudv1.TiDBMember, extraGroup.Spec.TiDB.Replicas)
			if extraGroup.Spec.TiDB.InstanceType != "" {
				for i := range members {
					members[i].InstanceType = extraGroup.Spec.TiDB.InstanceType
				}
			}
			extraGroup.Status.TiDB = &cloudv1.TiDBStatus{
				Status:  cloudv1.ComponentNormal,
				Members: members,
			}
			cr.Status.ExtraEndpoints = append(cr.Status.ExtraEndpoints, *extraGroup.Status)
		}
	}

	return cr
}

func TestScaleCluster(t *testing.T) {
	ctx := context.TODO()
	tests := []struct {
		name               string
		existCluster       *cloudv1.Cluster
		existConfigRevison *cloudv1.ConfigRevision
		req                *apiv1.ScaleClusterRequest
		wantErr            bool
		expect             func(*WithT, *cloudv1.Cluster, client.Client)
	}{
		{
			name: "scale up and out the default tidb group",
			existCluster: buildTestClusterCR(existClusterId, func(cr *cloudv1.Cluster) {
				cr.Spec.Cluster.TiDB = cloudv1.TiDB{
					Replicas:     2,
					InstanceType: "2c4g",
				}
			}),
			req: &apiv1.ScaleClusterRequest{
				ClusterId: existClusterId,
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIDB,
					InstanceType:  "8c16g",
					Count:         3,
				},
				TidbGroupId: defaultLoadBalancerID,
			},
			expect: func(g *WithT, actual *cloudv1.Cluster, _ client.Client) {
				// The spec should NOT be changed directly, instead a workflow annotation should be created
				g.Expect(actual.Spec.Cluster.TiDB).To(Equal(cloudv1.TiDB{
					Replicas:     2,
					InstanceType: "2c4g",
				}))
				g.Expect(actual.Spec.ExtraGroups).To(BeNil())
				g.Expect(actual.Annotations).To(HaveKey(consts.AnnoKeyScaleClusterWorkflow))
			},
		},
		{
			name: "disable gke heterogeneous deploy and scale tikv instance type",
			existCluster: buildScaleTestGKECluster(func(cr *cloudv1.Cluster) {
				cr.Spec.Cluster.TiKV = cloudv1.TiKV{
					Replicas:     3,
					InstanceType: "2c4g",
				}
				updateScaleTestGKEHeterogeneousIntent(t, cr, apicommon.ComponentType_TIKV, "", "2c4g", "4c8g")
			}),
			req: &apiv1.ScaleClusterRequest{
				ClusterId: existClusterId,
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIKV,
					InstanceType:  "4c8g",
					Count:         3,
				},
				TidbGroupId: "0",
				Options: map[string]string{
					scaleOptionDisableGKEHeterogeneous: "true",
				},
			},
			expect: func(g *WithT, actual *cloudv1.Cluster, _ client.Client) {
				g.Expect(actual.Spec.Cluster.TiKV.Replicas).To(Equal(3))
				g.Expect(actual.Spec.Cluster.TiKV.InstanceType).To(Equal("2c4g"))
				g.Expect(actual.Annotations).To(HaveKey(cloudv1.AnnoKeyGKEComputeClassForTiKVBindings))
				g.Expect(actual.Annotations).To(HaveKey(consts.AnnoKeyScaleClusterWorkflow))
			},
		},
		{
			name: "scale out the default tidb group without tidbgroup id",
			existCluster: buildTestClusterCR(existClusterId, func(cr *cloudv1.Cluster) {
				cr.Spec.Cluster.TiDB = cloudv1.TiDB{
					Replicas: 2,
				}
			}),
			req: &apiv1.ScaleClusterRequest{
				ClusterId: existClusterId,
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIDB,
					Count:         3,
				},
			},
			wantErr: true,
		},
		{
			name: "scale in and down the extra tidb group",
			existCluster: buildTestClusterCR(existClusterId, func(cr *cloudv1.Cluster) {
				cr.Spec.Cluster.TiDB = cloudv1.TiDB{
					Replicas:     3,
					InstanceType: "2c4g",
				}
				cr.Spec.ExtraGroups = []cloudv1.ExtraGroupSpec{
					{
						Name: "extra-1",
						TiDB: cloudv1.TiDB{
							Replicas:     1,
							InstanceType: "4c8g",
						},
					},
					{
						Name: "extra-2",
						TiDB: cloudv1.TiDB{
							Replicas:     2,
							InstanceType: "8c16g",
						},
					},
				}
			}),
			req: &apiv1.ScaleClusterRequest{
				ClusterId: existClusterId,
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIDB,
					InstanceType:  "2c4g",
					Count:         1,
				},
				TidbGroupId: "extra-2",
			},
			expect: func(g *WithT, actual *cloudv1.Cluster, _ client.Client) {
				// The spec should NOT be changed directly, instead a workflow annotation should be created
				g.Expect(actual.Spec.Cluster.TiDB).To(Equal(cloudv1.TiDB{
					Replicas:     3,
					InstanceType: "2c4g",
				}))
				g.Expect(actual.Spec.ExtraGroups).To(Equal([]cloudv1.ExtraGroupSpec{
					{
						Name: "extra-1",
						TiDB: cloudv1.TiDB{
							Replicas:     1,
							InstanceType: "4c8g",
						},
					},
					{
						Name: "extra-2",
						TiDB: cloudv1.TiDB{
							Replicas:     2,
							InstanceType: "8c16g",
						},
					},
				}))
				g.Expect(actual.Annotations).To(HaveKey(consts.AnnoKeyScaleClusterWorkflow))
			},
		},
		{
			name: "scale tiflash from 0",
			existCluster: buildTestClusterCR(existClusterId, func(cr *cloudv1.Cluster) {
				cr.Spec.Cluster.TiFlash = cloudv1.TiFlash{
					Replicas: 0,
				}
				cr.Spec.ConfigRevisionRef = &cloudv1.Reference{
					Name: "test-cr",
				}
			}),
			existConfigRevison: &cloudv1.ConfigRevision{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-cr",
				},
				Spec: cloudv1.ConfigRevisionSpec{
					Version: "v7.5.1",
				},
			},
			req: &apiv1.ScaleClusterRequest{
				ClusterId: existClusterId,
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIFLASH,
					InstanceType:  "2c4g",
					Count:         2,
					Storage: &apiv1.Cluster_Storage{
						Type:          "gp3",
						SizeGi:        100,
						Iops:          3000,
						ThroughputMbs: 125,
					},
				},
				ComponentConfigurations: []*apiv1.ComponentConfiguration{
					{
						ComponentType:             apicommon.ComponentType_TIFLASH,
						ConfigurationFileContents: `foo = "bar"`,
					},
				},
				ImageRepo: "tiflash",
				ImageTag:  "latest",
			},
			expect: func(g *WithT, actual *cloudv1.Cluster, cli client.Client) {

				g.Expect(actual.Annotations).To(HaveKey(consts.AnnoKeyScaleClusterWorkflow))
				// The spec should NOT be changed directly, instead a workflow annotation should be created
				g.Expect(actual.Spec.Cluster.TiFlash.Replicas).To(Equal(0))
				g.Expect(actual.Spec.Cluster.TiFlash.InstanceType).To(Equal(""))
				g.Expect(actual.Spec.Cluster.TiFlash.StorageClassName).To(Equal(""))
				g.Expect(actual.Spec.Cluster.TiFlash.StorageSize.IsZero()).To(BeTrue())
				var cr cloudv1.ConfigRevision
				g.Expect(actual.Spec.ConfigRevisionRef).ShouldNot(BeNil())
				g.Expect(cli.Get(ctx, client.ObjectKey{Name: actual.Spec.ConfigRevisionRef.Name}, &cr)).To(BeNil())
				g.Expect(cr.Spec.TiFlash).Should(BeNil())

			},
		},
		{
			name: "scale tikv raft volume",
			existCluster: buildTestClusterCR(existClusterId, func(cr *cloudv1.Cluster) {
				cr.Spec.Cluster.TiKV = cloudv1.TiKV{
					Replicas: 3,
					StorageVolumes: []cloudv1.ComponentStorageVolume{
						{
							Name:             "raft",
							StorageClassName: "tikv-gp3-3000-125",
							StorageSize:      resource.MustParse("100Gi"),
							MountPath:        "/var/lib/tikv/raft",
						},
					},
				}
				cr.Status.DB.TiKV = &cloudv1.TiKVStatus{
					StorageVolumes: []cloudv1.StorageVolumeStatus{
						{
							Name: "raft",
							VolumeStatus: cloudv1.VolumeStatus{
								BoundCount:           3,
								CurrentCount:         3,
								ModifiedCount:        3,
								CurrentCapacity:      resource.MustParse("100Gi"),
								ModifiedCapacity:     resource.MustParse("100Gi"),
								CurrentStorageClass:  "tikv-gp3-3000-125",
								ModifiedStorageClass: "tikv-gp3-3000-125",
								ResizedCount:         0,
								ResizedCapacity:      resource.MustParse("100Gi"),
							},
						},
					},
				}
			}),
			req: &apiv1.ScaleClusterRequest{
				ClusterId: existClusterId,
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIKV,
					Count:         3,
					AdditionalStorage: []*apiv1.Cluster_AdditionalStorage{
						{
							Name:          "raft",
							Type:          "io2",
							SizeGi:        200,
							Iops:          6000,
							ThroughputMbs: 250,
							MountPath:     "/var/lib/tikv/raft",
						},
					},
				},
				TidbGroupId: defaultLoadBalancerID, // TiKV uses the default group ID
			},
			expect: func(g *WithT, actual *cloudv1.Cluster, _ client.Client) {
				// The spec should NOT be changed directly, instead a workflow annotation should be created
				g.Expect(actual.Spec.Cluster.TiKV.StorageVolumes).To(Equal([]cloudv1.ComponentStorageVolume{
					{
						Name:             "raft",
						StorageClassName: "tikv-gp3-3000-125",
						StorageSize:      resource.MustParse("100Gi"),
						MountPath:        "/var/lib/tikv/raft",
					},
				}))
				g.Expect(actual.Annotations).To(HaveKey(consts.AnnoKeyScaleClusterWorkflow))
			},
		},
		{
			name: "scale tiproxy from 0 in default group",
			existCluster: buildTestClusterCR(existClusterId, func(cr *cloudv1.Cluster) {
				cr.Spec.Cluster.TiProxy = cloudv1.TiProxy{
					Replicas: 0,
				}
				cr.Spec.ConfigRevisionRef = &cloudv1.Reference{
					Name: "test-cr",
				}
				cr.Annotations[consts.DefaultLoadBalancerIDAnnotationKey] = defaultLoadBalancerID
			}),
			existConfigRevison: &cloudv1.ConfigRevision{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-cr",
				},
				Spec: cloudv1.ConfigRevisionSpec{
					Version: "v7.5.1",
				},
			},
			req: &apiv1.ScaleClusterRequest{
				ClusterId: existClusterId,
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIPROXY,
					InstanceType:  "2c4g",
					Count:         2,
				},
				ComponentConfigurations: []*apiv1.ComponentConfiguration{
					{
						ComponentType:             apicommon.ComponentType_TIPROXY,
						ConfigurationFileContents: `foo = "bar"`,
					},
				},
				ImageRepo:   "tiproxy",
				ImageTag:    "latest",
				TidbGroupId: defaultLoadBalancerID,
			},
			expect: func(g *WithT, actual *cloudv1.Cluster, cli client.Client) {
				g.Expect(actual.Spec.Cluster.TiProxy.Replicas).To(Equal(2))
				g.Expect(actual.Spec.Cluster.TiProxy.InstanceType).To(Equal("2c4g"))

				var cr cloudv1.ConfigRevision
				g.Expect(actual.Spec.ConfigRevisionRef).ShouldNot(BeNil())
				g.Expect(cli.Get(ctx, client.ObjectKey{Name: actual.Spec.ConfigRevisionRef.Name}, &cr)).To(BeNil())
				g.Expect(cr.Spec.TiProxy).ShouldNot(BeNil())
				g.Expect(*cr.Spec.TiProxy.BaseImage).To(Equal("tiproxy"))
				g.Expect(*cr.Spec.TiProxy.Tag).To(Equal("latest"))
				g.Expect(cr.Spec.TiProxy.Config.MP).To(Equal(map[string]interface{}{
					"foo": "bar",
				}))
				g.Expect(actual.Annotations[consts.AnnoKeyTiProxyToggleWorkflow]).To(
					ContainSubstring(fmt.Sprintf("%s-%s", consts.WorkflowTypeEnableTiProxy, defaultLoadBalancerID)),
				)
			},
		},
		{
			name: "scale tiproxy from 0 in extra tidb group",
			existCluster: buildTestClusterCR(existClusterId, func(cr *cloudv1.Cluster) {
				cr.Spec.ExtraGroups = []cloudv1.ExtraGroupSpec{
					{
						Name: "extra-1",
						TiDB: cloudv1.TiDB{
							InstanceType: "2c4g",
							Replicas:     1,
						},
						TiProxy: cloudv1.TiProxy{
							Replicas: 0,
						},
					},
					{
						Name: "extra-2", // deleted group
					},
				}
				cr.Spec.ConfigRevisionRef = &cloudv1.Reference{
					Name: "test-cr",
				}
				cr.Annotations[consts.DefaultLoadBalancerIDAnnotationKey] = defaultLoadBalancerID
			}),
			existConfigRevison: &cloudv1.ConfigRevision{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-cr",
				},
				Spec: cloudv1.ConfigRevisionSpec{
					Version: "v7.5.1",
				},
			},
			req: &apiv1.ScaleClusterRequest{
				ClusterId: existClusterId,
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIPROXY,
					InstanceType:  "2c4g",
					Count:         2,
				},
				ComponentConfigurations: []*apiv1.ComponentConfiguration{
					{
						ComponentType:             apicommon.ComponentType_TIPROXY,
						ConfigurationFileContents: `foo = "bar"`,
					},
				},
				ImageRepo:   "tiproxy",
				ImageTag:    "latest",
				TidbGroupId: "extra-1",
			},
			expect: func(g *WithT, actual *cloudv1.Cluster, cli client.Client) {
				g.Expect(actual.Spec.ExtraGroups).To(Equal([]cloudv1.ExtraGroupSpec{
					{
						Name: "extra-1",
						TiDB: cloudv1.TiDB{
							InstanceType: "2c4g",
							Replicas:     1,
						},
						TiProxy: cloudv1.TiProxy{
							Replicas:     2,
							InstanceType: "2c4g",
						},
					},
					{
						Name: "extra-2",
					},
				}))

				var cr cloudv1.ConfigRevision
				g.Expect(actual.Spec.ConfigRevisionRef).ShouldNot(BeNil())
				g.Expect(cli.Get(ctx, client.ObjectKey{Name: actual.Spec.ConfigRevisionRef.Name}, &cr)).To(BeNil())
				g.Expect(cr.Spec.TiProxy).ShouldNot(BeNil())
				g.Expect(*cr.Spec.TiProxy.BaseImage).To(Equal("tiproxy"))
				g.Expect(*cr.Spec.TiProxy.Tag).To(Equal("latest"))
				g.Expect(cr.Spec.TiProxy.Config.MP).To(Equal(map[string]interface{}{
					"foo": "bar",
				}))
				g.Expect(actual.Annotations[consts.AnnoKeyTiProxyToggleWorkflow]).To(
					ContainSubstring(fmt.Sprintf("%s-%s", consts.WorkflowTypeEnableTiProxy, "extra-1")),
				)
			},
		},
		{
			name: "scaling tiproxy from 1 to 2 in default group shouldn't change image and tag",
			existCluster: buildTestClusterCR(existClusterId, func(cr *cloudv1.Cluster) {
				cr.Spec.Cluster.TiProxy = cloudv1.TiProxy{
					Replicas:     1,
					InstanceType: "2c4g",
				}
				cr.Status.DB.TiProxy = &cloudv1.TiProxyStatus{
					Status: cloudv1.ComponentNormal,
					Members: []cloudv1.TiProxyMember{
						{
							Name:         "db-tiproxy-1",
							Status:       "Active",
							InstanceType: "2c4g",
						},
					},
					UpdatedReplicas:        1,
					InstanceTypeScalePhase: "",
				}
				cr.Spec.ConfigRevisionRef = &cloudv1.Reference{
					Name: "test-cr",
				}
				cr.Annotations[consts.DefaultLoadBalancerIDAnnotationKey] = defaultLoadBalancerID
			}),
			existConfigRevison: &cloudv1.ConfigRevision{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-cr",
				},
				Spec: cloudv1.ConfigRevisionSpec{
					Version: "v7.5.1",
					TiProxy: &cloudv1.Configuration{
						BaseImage: lo.ToPtr("old-base"),
						Tag:       lo.ToPtr("old-tag"),
						Config:    cloudv1.NewTomlConfig(map[string]interface{}{"a": "b"}),
					},
				},
			},
			req: &apiv1.ScaleClusterRequest{
				ClusterId: existClusterId,
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIPROXY,
					InstanceType:  "4c8g",
					Count:         2,
				},
				ComponentConfigurations: []*apiv1.ComponentConfiguration{
					{
						ComponentType:             apicommon.ComponentType_TIPROXY,
						ConfigurationFileContents: `foo = "bar"`,
					},
				},
				ImageRepo:   "tiproxy", // can't be changed
				ImageTag:    "latest",  // can't be changed
				TidbGroupId: defaultLoadBalancerID,
			},
			expect: func(g *WithT, actual *cloudv1.Cluster, cli client.Client) {
				// The spec should NOT be changed directly, instead a workflow annotation should be created
				g.Expect(actual.Spec.Cluster.TiProxy.Replicas).To(Equal(1))
				g.Expect(actual.Spec.Cluster.TiProxy.InstanceType).To(Equal("2c4g"))

				// The ConfigRevisionRef should still point to the old ConfigRevision
				// because the workflow approach doesn't update the cluster spec immediately
				var cr cloudv1.ConfigRevision
				g.Expect(actual.Spec.ConfigRevisionRef).ShouldNot(BeNil())
				g.Expect(actual.Spec.ConfigRevisionRef.Name).To(Equal("test-cr")) // Still pointing to old CR
				g.Expect(cli.Get(ctx, client.ObjectKey{Name: actual.Spec.ConfigRevisionRef.Name}, &cr)).To(BeNil())
				g.Expect(cr.Spec.TiProxy).ShouldNot(BeNil())
				// The old ConfigRevision should remain unchanged
				g.Expect(*cr.Spec.TiProxy.BaseImage).To(Equal("old-base"))
				g.Expect(*cr.Spec.TiProxy.Tag).To(Equal("old-tag"))
				g.Expect(cr.Spec.TiProxy.Config.MP).To(Equal(map[string]interface{}{
					"a": "b",
				}))
				g.Expect(actual.Annotations).To(HaveKey(consts.AnnoKeyScaleClusterWorkflow))
			},
		},
		{
			name: "scaling tiproxy from 1 to 0 in default group should be ok",
			existCluster: buildTestClusterCR(existClusterId, func(cr *cloudv1.Cluster) {
				cr.Spec.Cluster.TiProxy = cloudv1.TiProxy{
					Replicas:     1,
					InstanceType: "2c4g",
				}
				cr.Status.DB.TiProxy = &cloudv1.TiProxyStatus{
					Status: cloudv1.ComponentNormal,
					Members: []cloudv1.TiProxyMember{
						{
							Name:         "db-tiproxy-1",
							Status:       "Active",
							InstanceType: "2c4g",
						},
					},
					UpdatedReplicas:        1,
					InstanceTypeScalePhase: "",
				}
				cr.Spec.ConfigRevisionRef = &cloudv1.Reference{
					Name: "test-cr",
				}
				cr.Annotations[consts.DefaultLoadBalancerIDAnnotationKey] = defaultLoadBalancerID
			}),
			existConfigRevison: &cloudv1.ConfigRevision{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-cr",
				},
				Spec: cloudv1.ConfigRevisionSpec{
					Version: "v7.5.1",
					TiProxy: &cloudv1.Configuration{
						BaseImage: lo.ToPtr("old-base"),
						Tag:       lo.ToPtr("old-tag"),
						Config:    cloudv1.NewTomlConfig(map[string]interface{}{"a": "b"}),
					},
				},
			},
			req: &apiv1.ScaleClusterRequest{
				ClusterId: existClusterId,
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIPROXY,
					InstanceType:  "2c4g",
					Count:         0,
				},
				TidbGroupId: defaultLoadBalancerID,
			},
			expect: func(g *WithT, actual *cloudv1.Cluster, cli client.Client) {
				// TiProxy's replicas should keep unchanged when scaling to 0. Instead,
				// a tiproxy toggle workflow should be created.
				g.Expect(actual.Spec.Cluster.TiProxy.Replicas).To(Equal(1))
				g.Expect(actual.Annotations[consts.AnnoKeyTiProxyToggleWorkflow]).To(
					ContainSubstring(fmt.Sprintf("%s-%s", consts.WorkflowTypeDisableTiProxy, defaultLoadBalancerID)),
				)
				g.Expect(actual.Spec.Cluster.TiProxy.InstanceType).To(Equal("2c4g"))
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewWithT(t)
			var objs []client.Object
			if tt.existCluster != nil {
				objs = append(objs, tt.existCluster)
			}
			if tt.existConfigRevison != nil {
				objs = append(objs, tt.existConfigRevison)
			}
			// Add instance type configmap for scale operations
			objs = append(objs, buildInstanceTypeConfigMap())
			KubeClient := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithObjects(objs...).Build()
			reader := util.LabelSelectorReader{}
			s := ClusterServer{
				instanceTypeProvider: NewInstanceTypeProvider(),
				KubeClient:           KubeClient,
				clusterReader:        &reader,
			}

			_, err := s.ScaleCluster(ctx, tt.req)
			if tt.wantErr {
				g.Expect(err).ToNot(BeNil())
				return
			}
			g.Expect(err).To(BeNil())

			actualCluster, err := reader.Get(ctx, KubeClient, existClusterId)
			g.Expect(err).To(BeNil())
			tt.expect(g, actualCluster, KubeClient)
		})
	}
}

func TestValidateScaleClusterRequest(t *testing.T) {
	tests := []struct {
		name    string
		req     *apiv1.ScaleClusterRequest
		cluster *cloudv1.Cluster
		wantErr bool
	}{
		{
			name: "valid request",
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIDB,
					Count:         3,
				},
				TidbGroupId: "test-group",
			},
			wantErr: false,
		},
		{
			name: "invalid request - TIFLASH_LEARNER",
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIFLASH_LEARNER,
					Count:         3,
				},
				TidbGroupId: "test-group",
			},
			wantErr: true,
		},
		{
			name: "invalid request - missing TidbGroupId",
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIDB,
					Count:         3,
				},
			},
			wantErr: true,
		},
		{
			name: "invalid request - scale tidbgroup to 0",
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIDB,
					Count:         0,
				},
				TidbGroupId: "test-group",
			},
			wantErr: true,
		},
		{
			name: "scale tiflash from 0 with all necessary configuration",
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIFLASH,
					Count:         1,
					InstanceType:  "2c4g",
					Storage:       &apiv1.Cluster_Storage{},
				},
				ComponentConfigurations: []*apiv1.ComponentConfiguration{
					{
						ComponentType:             apicommon.ComponentType_TIFLASH,
						ConfigurationFileContents: "",
					},
				},
				ImageRepo: "tiflash",
				ImageTag:  "latest",
			},
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Spec.Cluster.TiFlash.Replicas = 0
			}),
			wantErr: false,
		},
		{
			name: "scale tiflash from 0 without instance type",
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIFLASH,
					Count:         1,
					Storage:       &apiv1.Cluster_Storage{},
				},
				ComponentConfigurations: []*apiv1.ComponentConfiguration{
					{
						ComponentType:             apicommon.ComponentType_TIFLASH,
						ConfigurationFileContents: "",
					},
				},
				ImageRepo: "tiflash",
				ImageTag:  "latest",
			},
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Spec.Cluster.TiFlash.Replicas = 0
			}),
			wantErr: true,
		},
		{
			name: "scale tiflash from 0 without storage",
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIFLASH,
					Count:         1,
					InstanceType:  "2c4g",
				},
				ComponentConfigurations: []*apiv1.ComponentConfiguration{
					{
						ComponentType:             apicommon.ComponentType_TIFLASH,
						ConfigurationFileContents: "",
					},
				},
				ImageRepo: "tiflash",
				ImageTag:  "latest",
			},
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Spec.Cluster.TiFlash.Replicas = 0
			}),
			wantErr: true,
		},
		{
			name: "scale tiflash from 0 without image",
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIFLASH,
					Count:         1,
					InstanceType:  "2c4g",
					Storage:       &apiv1.Cluster_Storage{},
				},
				ComponentConfigurations: []*apiv1.ComponentConfiguration{
					{
						ComponentType:             apicommon.ComponentType_TIFLASH,
						ConfigurationFileContents: "",
					},
				},
			},
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Spec.Cluster.TiFlash.Replicas = 0
			}),
			wantErr: true,
		},
		{
			name: "scale tiflash from 0 without configuration",
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIFLASH,
					Count:         1,
					InstanceType:  "2c4g",
					Storage:       &apiv1.Cluster_Storage{},
				},
			},
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Spec.Cluster.TiFlash.Replicas = 0
			}),
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := validateScaleClusterRequest(context.Background(), tt.req, tt.cluster)
			if (err != nil) != tt.wantErr {
				t.Errorf("validateScaleClusterRequest() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if err != nil {
				statusErr, _ := status.FromError(err)
				if statusErr.Code() != codes.InvalidArgument {
					t.Errorf("validateScaleClusterRequest() error code = %v, want %v", statusErr.Code(), codes.InvalidArgument)
				}
			}
		})
	}
}

func TestScaleFromZero(t *testing.T) {
	tests := []struct {
		name    string
		cluster *cloudv1.Cluster
		comp    cloudv1.ClusterComponent
		req     *apiv1.ScaleClusterRequest
		want    bool
	}{
		{
			name: "scale tiflash from 0 to 1 should be true",
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Spec.Cluster.TiFlash.Replicas = 0
			}),
			comp: &cloudv1.TiFlash{Replicas: 0},
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIFLASH,
					Count:         1,
				},
			},
			want: true,
		},
		{
			name: "scale tiflash from 1 to 2 should be false",
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Spec.Cluster.TiFlash.Replicas = 1
			}),
			comp: &cloudv1.TiFlash{Replicas: 1},
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIFLASH,
					Count:         2,
				},
			},
			want: false,
		},
		{
			name: "scale tiproxy from 0 to 1 with no tiproxy in any group should be true",
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Spec.Cluster.TiProxy.Replicas = 0
				c.Spec.ExtraGroups = []cloudv1.ExtraGroupSpec{
					{
						Name:    "extra-group",
						TiProxy: cloudv1.TiProxy{Replicas: 0},
					},
				}
			}),
			comp: &cloudv1.TiProxy{Replicas: 0},
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIPROXY,
					Count:         1,
				},
			},
			want: true,
		},
		{
			name: "scale tiproxy from 0 to 1 in extra group with 1 tiproxy in default group should be false",
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Spec.Cluster.TiProxy.Replicas = 1
				c.Spec.ExtraGroups = []cloudv1.ExtraGroupSpec{
					{
						Name:    "extra-group",
						TiProxy: cloudv1.TiProxy{Replicas: 0},
					},
				}
			}),
			comp: &cloudv1.TiProxy{Replicas: 0},
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIPROXY,
					Count:         1,
				},
			},
			want: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewWithT(t)
			got := scaleFromZero(tt.cluster, tt.comp, tt.req)
			g.Expect(got).To(Equal(tt.want))
		})
	}
}

func TestClusterServer_applyScaleCluster(t *testing.T) {
	ctx := context.Background()
	type args struct {
		cluster           *cloudv1.Cluster
		newConfigRevision string
		req               *apiv1.ScaleClusterRequest
	}
	tests := map[string]struct {
		args      args
		setup     func(*cloudv1.Cluster)
		wantErr   bool
		assertion func(*cloudv1.Cluster)
	}{
		"success: scale up TiDB in default group": {
			args: args{
				cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiDB = cloudv1.TiDB{
						Replicas:     2,
						InstanceType: "2c4g",
					}
				}),
				newConfigRevision: "",
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIDB,
						InstanceType:  "4c8g",
						Count:         3,
					},
					TidbGroupId: defaultLoadBalancerID,
				},
			},
			wantErr: false,
			assertion: func(c *cloudv1.Cluster) {
				// Normal TiDB scale operation uses workflow, spec should NOT be updated
				assert.Equal(t, 2, c.Spec.Cluster.TiDB.Replicas)
				assert.Equal(t, "2c4g", c.Spec.Cluster.TiDB.InstanceType)
				// Should have workflow annotation
				assert.Contains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
			},
		},
		"success: scale out TiProxy from 0 triggers the 'toggle-tiproxy' workflow annotation": {
			args: args{
				cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiProxy = cloudv1.TiProxy{
						Replicas: 0,
					}
					c.Spec.ConfigRevisionRef = &cloudv1.Reference{Name: "test-cr"}
				}),
				newConfigRevision: "",
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIPROXY,
						InstanceType:  "2c4g",
						Count:         1,
					},
					TidbGroupId: defaultLoadBalancerID,
					ComponentConfigurations: []*apiv1.ComponentConfiguration{
						{
							ComponentType:             apicommon.ComponentType_TIPROXY,
							ConfigurationFileContents: `foo = "bar"`,
						},
					},
					ImageRepo: "tiproxy",
					ImageTag:  "latest",
				},
			},
			wantErr: false,
			assertion: func(c *cloudv1.Cluster) {
				// For enable operations, the spec is updated directly
				assert.Equal(t, 1, c.Spec.Cluster.TiProxy.Replicas)
				assert.Equal(t, "2c4g", c.Spec.Cluster.TiProxy.InstanceType)
				assert.Contains(t, c.Annotations, consts.AnnoKeyTiProxyToggleWorkflow)
			},
		},
		"success: scale in TiProxy to 0 triggers the 'toggle-tiproxy' workflow": {
			args: args{
				cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiProxy = cloudv1.TiProxy{
						Replicas:     1,
						InstanceType: "old-type",
					}
					c.Spec.ConfigRevisionRef = &cloudv1.Reference{Name: "test-cr"}
				}),
				newConfigRevision: "",
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIPROXY,
						InstanceType:  "4c8g",
						Count:         0,
					},
					TidbGroupId: defaultLoadBalancerID,
				},
			},
			wantErr: false,
			assertion: func(c *cloudv1.Cluster) {
				// replicas should NOT be changed if the tiproxy is being disabled
				assert.Equal(t, 1, c.Spec.Cluster.TiProxy.Replicas)
				assert.Equal(t, "4c8g", c.Spec.Cluster.TiProxy.InstanceType)
				assert.Contains(t, c.Annotations, consts.AnnoKeyTiProxyToggleWorkflow)
			},
		},
		"success: dry-run update for no toggle type": {
			args: args{
				cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiDB = cloudv1.TiDB{
						Replicas:     2,
						InstanceType: "2c4g",
					}
				}),
				newConfigRevision: "",
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIDB,
						InstanceType:  "4c8g",
						Count:         3,
					},
					TidbGroupId: defaultLoadBalancerID,
				},
			},
			wantErr: false,
			assertion: func(c *cloudv1.Cluster) {
				// With the workflow approach, the spec should NOT be updated immediately
				assert.Equal(t, 2, c.Spec.Cluster.TiDB.Replicas)
				assert.Equal(t, "2c4g", c.Spec.Cluster.TiDB.InstanceType)
				assert.Contains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
			},
		},
		"fail: invalid storage spec": {
			args: args{
				cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiFlash = cloudv1.TiFlash{
						Replicas:         1,
						InstanceType:     "old-flash-type",
						StorageClassName: "tiflash-gp3-2000-100",
						StorageSize:      resource.MustParse("100Gi"),
					}
				}),
				newConfigRevision: "",
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIFLASH,
						InstanceType:  "flash-type",
						Count:         2,
						Storage: &apiv1.Cluster_Storage{
							Type: "", // invalid type
						},
					},
				},
			},
			wantErr:   true,
			assertion: func(_ *cloudv1.Cluster) {},
		},
		"fail: adding new storage volume for TiKV": {
			args: args{
				cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas: 3,
						StorageVolumes: []cloudv1.ComponentStorageVolume{
							{
								Name:             "raft",
								StorageClassName: "tikv-gp3-3000-125",
								StorageSize:      resource.MustParse("100Gi"),
								MountPath:        "/var/lib/tikv/raft",
							},
						},
					}
					c.Status.DB.TiKV = &cloudv1.TiKVStatus{}
				}),
				newConfigRevision: "",
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						Count:         3,
						AdditionalStorage: []*apiv1.Cluster_AdditionalStorage{
							{
								Name:          "raft",
								Type:          "io2",
								SizeGi:        200,
								Iops:          6000,
								ThroughputMbs: 250,
								MountPath:     "/var/lib/tikv/raft",
							},
							{
								Name:          "newvol",
								Type:          "gp3",
								SizeGi:        100,
								Iops:          3000,
								ThroughputMbs: 125,
								MountPath:     "/var/lib/tikv/newvol",
							},
						},
					},
				},
			},
			wantErr:   true,
			assertion: func(_ *cloudv1.Cluster) {},
		},
	}

	for name, tt := range tests {
		t.Run(name, func(t *testing.T) {
			// Add instance type configmap for scale operations
			cli := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithObjects(tt.args.cluster, buildInstanceTypeConfigMap()).Build()
			s := &ClusterServer{
				instanceTypeProvider: NewInstanceTypeProvider(),
				KubeClient:           cli,
			}
			logger := zap.NewNop()
			err := s.applyScaleCluster(ctx, logger, tt.args.cluster, tt.args.newConfigRevision, tt.args.req)
			if tt.wantErr {
				assert.Error(t, err)
				return
			}
			assert.NoError(t, err)
			// Fetch the updated cluster from K8s for enable/disable operations
			var updatedCluster cloudv1.Cluster
			err = cli.Get(ctx, client.ObjectKey{Name: tt.args.cluster.Name}, &updatedCluster)
			assert.NoError(t, err)
			tt.assertion(&updatedCluster)
		})
	}
}

func TestClusterServer_applyScaleClusterDisableGKEHeterogeneous(t *testing.T) {
	ctx := context.Background()
	tikvIntentKey := cloudv1.AnnoKeyGKEComputeClassForTiKVBindings

	type args struct {
		cluster *cloudv1.Cluster
		req     *apiv1.ScaleClusterRequest
	}
	tests := map[string]struct {
		args      args
		wantCode  codes.Code
		assertion func(*cloudv1.Cluster)
	}{
		"intent exists and option true defers intent removal to instance type workflow step": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas:     3,
						InstanceType: "2c4g",
					}
					updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIKV, "", "2c4g", "4c8g")
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						InstanceType:  "4c8g",
						Count:         3,
					},
					TidbGroupId: "0",
					Options: map[string]string{
						scaleOptionDisableGKEHeterogeneous: "true",
					},
				},
			},
			wantCode: codes.OK,
			assertion: func(c *cloudv1.Cluster) {
				assert.Equal(t, "2c4g", c.Spec.Cluster.TiKV.InstanceType)
				assert.Contains(t, c.Annotations, tikvIntentKey)
				assert.Contains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
				assertScaleWorkflowDisableGKEHeterogeneous(t, c, true)
			},
		},
		"option true with unchanged target rejects count-only workflow": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas:     3,
						InstanceType: "2c4g",
					}
					updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIKV, "", "2c4g", "4c8g")
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						InstanceType:  "2c4g",
						Count:         4,
					},
					Options: map[string]string{
						scaleOptionDisableGKEHeterogeneous: "true",
					},
				},
			},
			wantCode: codes.InvalidArgument,
			assertion: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, tikvIntentKey)
				assert.NotContains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
			},
		},
		"option true with unchanged target rejects finalize-only workflow": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas:     3,
						InstanceType: "2c4g",
					}
					updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIKV, "", "2c4g", "4c8g")
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						InstanceType:  "2c4g",
						Count:         3,
					},
					Options: map[string]string{
						scaleOptionDisableGKEHeterogeneous: "true",
					},
				},
			},
			wantCode: codes.InvalidArgument,
			assertion: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, tikvIntentKey)
				assert.NotContains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
			},
		},
		"option true with empty target rejects before clearing intent": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas:     3,
						InstanceType: "2c4g",
					}
					updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIKV, "", "2c4g", "4c8g")
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						Count:         4,
					},
					ComponentConfigurations: []*apiv1.ComponentConfiguration{
						{
							ComponentType:             apicommon.ComponentType_TIKV,
							ConfigurationFileContents: `foo = "bar"`,
						},
					},
					Options: map[string]string{
						scaleOptionDisableGKEHeterogeneous: "true",
					},
				},
			},
			wantCode: codes.InvalidArgument,
			assertion: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, tikvIntentKey)
				assert.NotContains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
			},
		},
		"option absent rejects tikv instance type mutation while intent exists": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas:     3,
						InstanceType: "2c4g",
					}
					updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIKV, "", "2c4g", "4c8g")
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						InstanceType:  "4c8g",
						Count:         3,
					},
				},
			},
			wantCode: codes.FailedPrecondition,
			assertion: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, tikvIntentKey)
				assert.NotContains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
			},
		},
		"option absent allows tikv count-only scale while intent exists": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas:     3,
						InstanceType: "2c4g",
					}
					updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIKV, "", "2c4g", "4c8g")
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						Count:         4,
					},
				},
			},
			wantCode: codes.OK,
			assertion: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, tikvIntentKey)
				assert.Contains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
				assertScaleWorkflowDesiredTiKV(t, c, 4, "2c4g")
			},
		},
		"option absent rejects tidb group instance type mutation while intent exists": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiDB = cloudv1.TiDB{
						Replicas:     2,
						InstanceType: "2c4g",
					}
					updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIDB, defaultLoadBalancerID, "2c4g", "4c8g")
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIDB,
						InstanceType:  "4c8g",
						Count:         2,
					},
					TidbGroupId: defaultLoadBalancerID,
				},
			},
			wantCode: codes.FailedPrecondition,
			assertion: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, c.TiDBGKEComputeClassBindingsAnnoKey(defaultLoadBalancerID))
				assert.NotContains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
			},
		},
		"option absent allows tidb group count-only scale while intent exists": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiDB = cloudv1.TiDB{
						Replicas:     2,
						InstanceType: "2c4g",
					}
					updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIDB, defaultLoadBalancerID, "2c4g", "4c8g")
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIDB,
						Count:         3,
					},
					TidbGroupId: defaultLoadBalancerID,
				},
			},
			wantCode: codes.OK,
			assertion: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, c.TiDBGKEComputeClassBindingsAnnoKey(defaultLoadBalancerID))
				assert.Contains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
				assertScaleWorkflowDesiredTiDB(t, c, 3, "2c4g")
			},
		},
		"intent missing and target differs creates ordinary workflow": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas:     3,
						InstanceType: "2c4g",
					}
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						InstanceType:  "4c8g",
						Count:         4,
					},
					Options: map[string]string{
						scaleOptionDisableGKEHeterogeneous: "true",
					},
				},
			},
			wantCode: codes.OK,
			assertion: func(c *cloudv1.Cluster) {
				assert.NotContains(t, c.Annotations, tikvIntentKey)
				assert.Contains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
				assertScaleWorkflowDisableGKEHeterogeneous(t, c, false)
				assertScaleWorkflowDesiredTiKV(t, c, 4, "4c8g")
			},
		},
		"non gke ignores disable option": {
			args: args{
				cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
					c.Status.DB.TiKV = &cloudv1.TiKVStatus{}
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas:     3,
						InstanceType: "2c4g",
					}
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						InstanceType:  "4c8g",
						Count:         4,
					},
					Options: map[string]string{
						scaleOptionDisableGKEHeterogeneous: "true",
					},
				},
			},
			wantCode: codes.OK,
			assertion: func(c *cloudv1.Cluster) {
				assert.NotContains(t, c.Annotations, tikvIntentKey)
				assert.Contains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
				assertScaleWorkflowDisableGKEHeterogeneous(t, c, false)
				assertScaleWorkflowDesiredTiKV(t, c, 4, "4c8g")
			},
		},
		"intent missing and target already applied creates ordinary workflow": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas:     3,
						InstanceType: "4c8g",
					}
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						InstanceType:  "4c8g",
						Count:         3,
					},
					TidbGroupId: "0",
					Options: map[string]string{
						scaleOptionDisableGKEHeterogeneous: "true",
					},
				},
			},
			wantCode: codes.OK,
			assertion: func(c *cloudv1.Cluster) {
				assert.NotContains(t, c.Annotations, tikvIntentKey)
				assert.Contains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
				assertScaleWorkflowDisableGKEHeterogeneous(t, c, false)
			},
		},
		"option true allows combined replica change while deferring intent removal": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas:     3,
						InstanceType: "2c4g",
					}
					updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIKV, "", "2c4g", "4c8g")
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						InstanceType:  "4c8g",
						Count:         4,
					},
					Options: map[string]string{
						scaleOptionDisableGKEHeterogeneous: "true",
					},
				},
			},
			wantCode: codes.OK,
			assertion: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, tikvIntentKey)
				assert.Contains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
				assertScaleWorkflowDisableGKEHeterogeneous(t, c, true)
				assertScaleWorkflowDesiredTiKV(t, c, 4, "4c8g")
			},
		},
		"option true allows storage change while deferring intent removal": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas:         3,
						InstanceType:     "2c4g",
						StorageClassName: "tikv-gp3-3000-125",
						StorageSize:      resource.MustParse("100Gi"),
					}
					updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIKV, "", "2c4g", "4c8g")
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						InstanceType:  "4c8g",
						Count:         3,
						Storage: &apiv1.Cluster_Storage{
							Type:          "io2",
							SizeGi:        200,
							Iops:          6000,
							ThroughputMbs: 250,
						},
					},
					Options: map[string]string{
						scaleOptionDisableGKEHeterogeneous: "true",
					},
				},
			},
			wantCode: codes.OK,
			assertion: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, tikvIntentKey)
				assert.Contains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
				assertScaleWorkflowDisableGKEHeterogeneous(t, c, true)
				desired := scaleWorkflowDesiredTiKV(t, c)
				assert.Equal(t, 3, desired.Replicas)
				assert.Equal(t, "4c8g", desired.InstanceType)
				assert.Equal(t, "tikv-io2-6000-250", desired.StorageClassName)
				assert.Equal(t, resource.MustParse("200Gi"), desired.StorageSize)
			},
		},
		"option true allows additional storage change while deferring intent removal": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas:     3,
						InstanceType: "2c4g",
						StorageVolumes: []cloudv1.ComponentStorageVolume{
							{
								Name:             "raft",
								StorageClassName: "tikv-gp3-3000-125",
								StorageSize:      resource.MustParse("100Gi"),
								MountPath:        "/var/lib/tikv/raft",
							},
						},
					}
					c.Status.DB.TiKV = &cloudv1.TiKVStatus{
						StorageVolumes: []cloudv1.StorageVolumeStatus{
							{
								Name: "raft",
								VolumeStatus: cloudv1.VolumeStatus{
									BoundCount:           3,
									CurrentCount:         3,
									ModifiedCount:        3,
									CurrentCapacity:      resource.MustParse("100Gi"),
									ModifiedCapacity:     resource.MustParse("100Gi"),
									CurrentStorageClass:  "tikv-gp3-3000-125",
									ModifiedStorageClass: "tikv-gp3-3000-125",
									ResizedCount:         0,
									ResizedCapacity:      resource.MustParse("100Gi"),
								},
							},
						},
					}
					updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIKV, "", "2c4g", "4c8g")
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						InstanceType:  "4c8g",
						Count:         3,
						AdditionalStorage: []*apiv1.Cluster_AdditionalStorage{
							{
								Name:          "raft",
								Type:          "io2",
								SizeGi:        200,
								Iops:          6000,
								ThroughputMbs: 250,
								MountPath:     "/var/lib/tikv/raft",
							},
						},
					},
					Options: map[string]string{
						scaleOptionDisableGKEHeterogeneous: "true",
					},
				},
			},
			wantCode: codes.OK,
			assertion: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, tikvIntentKey)
				assert.Contains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
				assertScaleWorkflowDisableGKEHeterogeneous(t, c, true)
				desired := scaleWorkflowDesiredTiKV(t, c)
				assert.Equal(t, 3, desired.Replicas)
				assert.Equal(t, "4c8g", desired.InstanceType)
				assert.Len(t, desired.StorageVolumes, 1)
				assert.Equal(t, "raft", desired.StorageVolumes[0].Name)
				assert.Equal(t, "tikv-io2-6000-250", desired.StorageVolumes[0].StorageClassName)
				assert.True(t, desired.StorageVolumes[0].StorageSize.Equal(storageSizeGi2quantity(200)))
				assert.Equal(t, "/var/lib/tikv/raft", desired.StorageVolumes[0].MountPath)
			},
		},
		"option true defers scale to zero and intent removal to workflow": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas:     3,
						InstanceType: "2c4g",
					}
					updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIKV, "", "2c4g", "4c8g")
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						InstanceType:  "4c8g",
						Count:         0,
					},
					Options: map[string]string{
						scaleOptionDisableGKEHeterogeneous: "true",
					},
				},
			},
			wantCode: codes.OK,
			assertion: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, tikvIntentKey)
				assert.Contains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
				assert.Equal(t, 3, c.Spec.Cluster.TiKV.Replicas)
				assert.Equal(t, "2c4g", c.Spec.Cluster.TiKV.InstanceType)
				assertScaleWorkflowDisableGKEHeterogeneous(t, c, true)
				assertScaleWorkflowDesiredTiKV(t, c, 0, "4c8g")
			},
		},
		"tidb default group intent exists and option true defers group intent removal": {
			args: args{
				cluster: buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiDB = cloudv1.TiDB{
						Replicas:     2,
						InstanceType: "2c4g",
					}
					updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIDB, defaultLoadBalancerID, "2c4g", "4c8g")
				}),
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIDB,
						InstanceType:  "4c8g",
						Count:         3,
					},
					TidbGroupId: defaultLoadBalancerID,
					Options: map[string]string{
						scaleOptionDisableGKEHeterogeneous: "true",
					},
				},
			},
			wantCode: codes.OK,
			assertion: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, c.TiDBGKEComputeClassBindingsAnnoKey(defaultLoadBalancerID))
				assert.Contains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
				assertScaleWorkflowDisableGKEHeterogeneous(t, c, true)
				assertScaleWorkflowDesiredTiDB(t, c, 3, "4c8g")
			},
		},
	}

	for name, tt := range tests {
		t.Run(name, func(t *testing.T) {
			cli := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithObjects(tt.args.cluster, buildInstanceTypeConfigMap()).Build()
			s := &ClusterServer{
				instanceTypeProvider: NewInstanceTypeProvider(),
				KubeClient:           cli,
			}
			err := s.applyScaleCluster(ctx, zap.NewNop(), tt.args.cluster, tt.args.cluster.Spec.ConfigRevisionRef.Name, tt.args.req)
			if tt.wantCode != codes.OK {
				assert.Error(t, err)
				assert.Equal(t, tt.wantCode, status.Code(err))
			} else {
				assert.NoError(t, err)
			}

			var updatedCluster cloudv1.Cluster
			assert.NoError(t, cli.Get(ctx, client.ObjectKey{Name: tt.args.cluster.Name}, &updatedCluster))
			tt.assertion(&updatedCluster)
		})
	}
}

func assertScaleWorkflowDisableGKEHeterogeneous(t *testing.T, cluster *cloudv1.Cluster, want bool) {
	t.Helper()
	wf := scaleWorkflowFromCluster(t, cluster)
	if want {
		assert.Equal(t, "true", wf.Spec.Inputs[consts.WfInputKeyScaleClusterDisableGKEHeterogeneous])
		return
	}
	assert.NotContains(t, wf.Spec.Inputs, consts.WfInputKeyScaleClusterDisableGKEHeterogeneous)
}

func scaleWorkflowFromCluster(t *testing.T, cluster *cloudv1.Cluster) cloudv1.Workflow {
	t.Helper()
	var wf cloudv1.Workflow
	assert.NoError(t, yaml.Unmarshal([]byte(cluster.Annotations[consts.AnnoKeyScaleClusterWorkflow]), &wf))
	return wf
}

func scaleWorkflowDesiredTiKV(t *testing.T, cluster *cloudv1.Cluster) cloudv1.TiKV {
	t.Helper()
	wf := scaleWorkflowFromCluster(t, cluster)
	var desired cloudv1.TiKV
	assert.NoError(t, yaml.Unmarshal([]byte(wf.Spec.Inputs[consts.WfInputKeyScaleClusterDesiredComponentSpec]), &desired))
	return desired
}

func assertScaleWorkflowDesiredTiKV(t *testing.T, cluster *cloudv1.Cluster, wantReplicas int, wantInstanceType string) {
	t.Helper()
	desired := scaleWorkflowDesiredTiKV(t, cluster)
	assert.Equal(t, wantReplicas, desired.Replicas)
	assert.Equal(t, wantInstanceType, desired.InstanceType)
}

func assertScaleWorkflowDesiredTiDB(t *testing.T, cluster *cloudv1.Cluster, wantReplicas int, wantInstanceType string) {
	t.Helper()
	wf := scaleWorkflowFromCluster(t, cluster)
	var desired cloudv1.TiDB
	assert.NoError(t, yaml.Unmarshal([]byte(wf.Spec.Inputs[consts.WfInputKeyScaleClusterDesiredComponentSpec]), &desired))
	assert.Equal(t, wantReplicas, desired.Replicas)
	assert.Equal(t, wantInstanceType, desired.InstanceType)
}

func TestScaleClusterDisableGKEHeterogeneousAllowsConfigurationWithInstanceTypeChange(t *testing.T) {
	ctx := context.Background()
	cluster := buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
		c.Spec.Cluster.TiKV = cloudv1.TiKV{
			Replicas:     3,
			InstanceType: "2c4g",
		}
		updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIKV, "", "2c4g", "4c8g")
	})
	existingConfigRevision := &cloudv1.ConfigRevision{
		ObjectMeta: metav1.ObjectMeta{
			Name: cluster.Spec.ConfigRevisionRef.Name,
		},
		Spec: cloudv1.ConfigRevisionSpec{
			TiKV: &cloudv1.Configuration{
				Config: cloudv1.NewTomlConfig(map[string]interface{}{"tikv": "b"}),
			},
			Version: "v7.5.1",
		},
	}
	cli := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithObjects(cluster, existingConfigRevision, buildInstanceTypeConfigMap()).Build()
	reader := util.LabelSelectorReader{}
	s := &ClusterServer{
		instanceTypeProvider: NewInstanceTypeProvider(),
		KubeClient:           cli,
		clusterReader:        &reader,
	}

	_, err := s.ScaleCluster(ctx, &apiv1.ScaleClusterRequest{
		ClusterId: existClusterId,
		NodeSpec: &apiv1.Cluster_NodeSpec{
			ComponentType: apicommon.ComponentType_TIKV,
			InstanceType:  "4c8g",
			Count:         3,
		},
		ComponentConfigurations: []*apiv1.ComponentConfiguration{
			{
				ComponentType:             apicommon.ComponentType_TIKV,
				ConfigurationFileContents: `tikv = "c"`,
			},
		},
		Options: map[string]string{
			scaleOptionDisableGKEHeterogeneous: "true",
		},
	})
	assert.NoError(t, err)

	var updatedCluster cloudv1.Cluster
	assert.NoError(t, cli.Get(ctx, client.ObjectKey{Name: cluster.Name}, &updatedCluster))
	assert.Contains(t, updatedCluster.Annotations, cloudv1.AnnoKeyGKEComputeClassForTiKVBindings)
	assert.Contains(t, updatedCluster.Annotations, consts.AnnoKeyScaleClusterWorkflow)
	assertScaleWorkflowDisableGKEHeterogeneous(t, &updatedCluster, true)
	assertScaleWorkflowDesiredTiKV(t, &updatedCluster, 3, "4c8g")

	wf := scaleWorkflowFromCluster(t, &updatedCluster)
	desiredConfigRevisionName := wf.Spec.Inputs[consts.WfInputKeyScaleClusterDesiredConfigRevisionName]
	assert.NotEmpty(t, desiredConfigRevisionName)
	assert.NotEqual(t, existingConfigRevision.Name, desiredConfigRevisionName)

	var desiredConfigRevision cloudv1.ConfigRevision
	assert.NoError(t, cli.Get(ctx, client.ObjectKey{Name: desiredConfigRevisionName}, &desiredConfigRevision))
	assert.Equal(t, "c", desiredConfigRevision.Spec.TiKV.Config.MP["tikv"])
}

func TestScaleClusterDisableGKEHeterogeneousIgnoredWithoutIntent(t *testing.T) {
	ctx := context.Background()
	cluster := buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
		c.Spec.Cluster.TiKV = cloudv1.TiKV{
			Replicas:     3,
			InstanceType: "2c4g",
		}
	})
	existingConfigRevision := &cloudv1.ConfigRevision{
		ObjectMeta: metav1.ObjectMeta{
			Name: cluster.Spec.ConfigRevisionRef.Name,
		},
		Spec: cloudv1.ConfigRevisionSpec{
			TiKV: &cloudv1.Configuration{
				Config: cloudv1.NewTomlConfig(map[string]interface{}{"tikv": "b"}),
			},
			Version: "v7.5.1",
		},
	}
	cli := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithObjects(cluster, existingConfigRevision, buildInstanceTypeConfigMap()).Build()
	reader := util.LabelSelectorReader{}
	s := &ClusterServer{
		instanceTypeProvider: NewInstanceTypeProvider(),
		KubeClient:           cli,
		clusterReader:        &reader,
	}

	_, err := s.ScaleCluster(ctx, &apiv1.ScaleClusterRequest{
		ClusterId: existClusterId,
		NodeSpec: &apiv1.Cluster_NodeSpec{
			ComponentType: apicommon.ComponentType_TIKV,
			InstanceType:  "4c8g",
			Count:         4,
		},
		ComponentConfigurations: []*apiv1.ComponentConfiguration{
			{
				ComponentType:             apicommon.ComponentType_TIKV,
				ConfigurationFileContents: `tikv = "c"`,
			},
		},
		Options: map[string]string{
			scaleOptionDisableGKEHeterogeneous: "true",
		},
	})
	assert.NoError(t, err)

	var updatedCluster cloudv1.Cluster
	assert.NoError(t, cli.Get(ctx, client.ObjectKey{Name: cluster.Name}, &updatedCluster))
	assert.NotContains(t, updatedCluster.Annotations, cloudv1.AnnoKeyGKEComputeClassForTiKVBindings)
	assert.Contains(t, updatedCluster.Annotations, consts.AnnoKeyScaleClusterWorkflow)
	assertScaleWorkflowDisableGKEHeterogeneous(t, &updatedCluster, false)
	assertScaleWorkflowDesiredTiKV(t, &updatedCluster, 4, "4c8g")

	wf := scaleWorkflowFromCluster(t, &updatedCluster)
	desiredConfigRevisionName := wf.Spec.Inputs[consts.WfInputKeyScaleClusterDesiredConfigRevisionName]
	assert.NotEmpty(t, desiredConfigRevisionName)
	assert.NotEqual(t, existingConfigRevision.Name, desiredConfigRevisionName)
}

func TestScaleClusterGKEHeterogeneousEmptyInstanceTypeWithConfigurationRejectsBeforeConfigUpdate(t *testing.T) {
	ctx := context.Background()
	cluster := buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
		c.Spec.Cluster.TiKV = cloudv1.TiKV{
			Replicas:     3,
			InstanceType: "2c4g",
		}
		updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIKV, "", "2c4g", "4c8g")
	})
	existingConfigRevision := &cloudv1.ConfigRevision{
		ObjectMeta: metav1.ObjectMeta{
			Name: cluster.Spec.ConfigRevisionRef.Name,
		},
		Spec: cloudv1.ConfigRevisionSpec{
			Version: "v7.5.1",
		},
	}
	cli := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithObjects(cluster, existingConfigRevision, buildInstanceTypeConfigMap()).Build()
	reader := util.LabelSelectorReader{}
	s := &ClusterServer{
		instanceTypeProvider: NewInstanceTypeProvider(),
		KubeClient:           cli,
		clusterReader:        &reader,
	}

	_, err := s.ScaleCluster(ctx, &apiv1.ScaleClusterRequest{
		ClusterId: existClusterId,
		NodeSpec: &apiv1.Cluster_NodeSpec{
			ComponentType: apicommon.ComponentType_TIKV,
			Count:         4,
		},
		ComponentConfigurations: []*apiv1.ComponentConfiguration{
			{
				ComponentType:             apicommon.ComponentType_TIKV,
				ConfigurationFileContents: `foo = "bar"`,
			},
		},
	})
	assert.Error(t, err)
	assert.Equal(t, codes.InvalidArgument, status.Code(err))

	var updatedCluster cloudv1.Cluster
	assert.NoError(t, cli.Get(ctx, client.ObjectKey{Name: cluster.Name}, &updatedCluster))
	assert.Contains(t, updatedCluster.Annotations, cloudv1.AnnoKeyGKEComputeClassForTiKVBindings)
	assert.NotContains(t, updatedCluster.Annotations, consts.AnnoKeyScaleClusterWorkflow)

	var configRevisions cloudv1.ConfigRevisionList
	assert.NoError(t, cli.List(ctx, &configRevisions))
	assert.Len(t, configRevisions.Items, 1)
	assert.Equal(t, existingConfigRevision.Name, configRevisions.Items[0].Name)
}

func TestScaleClusterGKEHeterogeneousInstanceTypeMutationRejectsBeforeConfigUpdate(t *testing.T) {
	ctx := context.Background()
	cluster := buildScaleTestGKECluster(func(c *cloudv1.Cluster) {
		c.Spec.Cluster.TiKV = cloudv1.TiKV{
			Replicas:     3,
			InstanceType: "2c4g",
		}
		updateScaleTestGKEHeterogeneousIntent(t, c, apicommon.ComponentType_TIKV, "", "2c4g", "4c8g")
	})
	existingConfigRevision := &cloudv1.ConfigRevision{
		ObjectMeta: metav1.ObjectMeta{
			Name: cluster.Spec.ConfigRevisionRef.Name,
		},
		Spec: cloudv1.ConfigRevisionSpec{
			Version: "v7.5.1",
		},
	}
	cli := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithObjects(cluster, existingConfigRevision, buildInstanceTypeConfigMap()).Build()
	reader := util.LabelSelectorReader{}
	s := &ClusterServer{
		instanceTypeProvider: NewInstanceTypeProvider(),
		KubeClient:           cli,
		clusterReader:        &reader,
	}

	_, err := s.ScaleCluster(ctx, &apiv1.ScaleClusterRequest{
		ClusterId: existClusterId,
		NodeSpec: &apiv1.Cluster_NodeSpec{
			ComponentType: apicommon.ComponentType_TIKV,
			InstanceType:  "4c8g",
			Count:         3,
		},
		ComponentConfigurations: []*apiv1.ComponentConfiguration{
			{
				ComponentType:             apicommon.ComponentType_TIKV,
				ConfigurationFileContents: `foo = "bar"`,
			},
		},
	})
	assert.Error(t, err)
	assert.Equal(t, codes.FailedPrecondition, status.Code(err))

	var updatedCluster cloudv1.Cluster
	assert.NoError(t, cli.Get(ctx, client.ObjectKey{Name: cluster.Name}, &updatedCluster))
	assert.Contains(t, updatedCluster.Annotations, cloudv1.AnnoKeyGKEComputeClassForTiKVBindings)
	assert.NotContains(t, updatedCluster.Annotations, consts.AnnoKeyScaleClusterWorkflow)

	var configRevisions cloudv1.ConfigRevisionList
	assert.NoError(t, cli.List(ctx, &configRevisions))
	assert.Len(t, configRevisions.Items, 1)
	assert.Equal(t, existingConfigRevision.Name, configRevisions.Items[0].Name)
}

func buildScaleTestGKECluster(apply func(*cloudv1.Cluster)) *cloudv1.Cluster {
	return buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
		c.Labels[cloudv1.ProviderTypeLabel] = cloudv1.ProviderTypeGKE
		c.Status.DB.TiKV = &cloudv1.TiKVStatus{}
		if apply != nil {
			apply(c)
		}
	})
}

func updateScaleTestGKEHeterogeneousIntent(t *testing.T, cluster *cloudv1.Cluster, componentType apicommon.ComponentType, groupID, baseInstanceType, extraInstanceType string) {
	t.Helper()
	key := cloudv1.AnnoKeyGKEComputeClassForTiKVBindings
	if componentType == apicommon.ComponentType_TIDB {
		key = cluster.TiDBGKEComputeClassBindingsAnnoKey(groupID)
	}
	updateTestGKEComputeClassBinding(t, cluster, key, "", []string{baseInstanceType, extraInstanceType}, []string{"base-pool", "extra-pool"})
}

func Test_applyClusterScaleChanges(t *testing.T) {
	ctx := context.Background()
	type args struct {
		cluster           *cloudv1.Cluster
		newConfigRevision string
		req               *apiv1.ScaleClusterRequest
	}
	tests := map[string]struct {
		args               args
		wantErr            bool
		wantCompToggleType string
		assertion          func(*cloudv1.Cluster)
	}{
		"update replicas and instance type for TiDB": {
			args: args{
				cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiDB = cloudv1.TiDB{
						Replicas:     2,
						InstanceType: "2c4g",
					}
				}),
				newConfigRevision: "",
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIDB,
						InstanceType:  "4c8g",
						Count:         3,
					},
					TidbGroupId: defaultLoadBalancerID,
				},
			},
			wantErr: false,
			assertion: func(c *cloudv1.Cluster) {
				assert.Equal(t, 3, c.Spec.Cluster.TiDB.Replicas)
				assert.Equal(t, "4c8g", c.Spec.Cluster.TiDB.InstanceType)
			},
		},
		"update storage for TiFlash": {
			args: args{
				cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiFlash = cloudv1.TiFlash{
						Replicas:         1,
						InstanceType:     "old-flash-type",
						StorageClassName: "tiflash-gp3-2000-100",
						StorageSize:      resource.MustParse("100Gi"),
					}
				}),
				newConfigRevision: "",
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIFLASH,
						InstanceType:  "flash-type",
						Count:         2,
						Storage: &apiv1.Cluster_Storage{
							Type:          "io2",
							SizeGi:        200,
							Iops:          3000,
							ThroughputMbs: 125,
						},
					},
				},
			},
			wantErr: false,
			assertion: func(c *cloudv1.Cluster) {
				assert.Equal(t, 2, c.Spec.Cluster.TiFlash.Replicas)
				assert.Equal(t, "flash-type", c.Spec.Cluster.TiFlash.InstanceType)
				assert.Equal(t, "tiflash-io2-3000-125", c.Spec.Cluster.TiFlash.StorageClassName)
				assert.Equal(t, resource.MustParse("200Gi"), c.Spec.Cluster.TiFlash.StorageSize)
			},
		},
		"update additional storage for TiKV": {
			args: args{
				cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas: 3,
						StorageVolumes: []cloudv1.ComponentStorageVolume{
							{
								Name:             "raft",
								StorageClassName: "tikv-gp3-3000-125",
								StorageSize:      resource.MustParse("100Gi"),
								MountPath:        "/var/lib/tikv/raft",
							},
						},
					}
					c.Status.DB.TiKV = &cloudv1.TiKVStatus{}
				}),
				newConfigRevision: "",
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						Count:         3,
						AdditionalStorage: []*apiv1.Cluster_AdditionalStorage{
							{
								Name:          "raft",
								Type:          "io2",
								SizeGi:        200,
								Iops:          6000,
								ThroughputMbs: 250,
								MountPath:     "/var/lib/tikv/raft",
							},
						},
					},
				},
			},
			wantErr: false,
			assertion: func(c *cloudv1.Cluster) {
				assert.Equal(t, []cloudv1.ComponentStorageVolume{
					{
						Name:             "raft",
						StorageClassName: "tikv-io2-6000-250",
						StorageSize:      storageSizeGi2quantity(200),
						MountPath:        "/var/lib/tikv/raft",
					},
				}, c.Spec.Cluster.TiKV.StorageVolumes)
			},
		},
		"fail on adding new storage volume for TiKV": {
			args: args{
				cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiKV = cloudv1.TiKV{
						Replicas: 3,
						StorageVolumes: []cloudv1.ComponentStorageVolume{
							{
								Name:             "raft",
								StorageClassName: "tikv-gp3-3000-125",
								StorageSize:      resource.MustParse("100Gi"),
								MountPath:        "/var/lib/tikv/raft",
							},
						},
					}
					c.Status.DB.TiKV = &cloudv1.TiKVStatus{}
				}),
				newConfigRevision: "",
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIKV,
						Count:         3,
						AdditionalStorage: []*apiv1.Cluster_AdditionalStorage{
							{
								Name:          "raft",
								Type:          "io2",
								SizeGi:        200,
								Iops:          6000,
								ThroughputMbs: 250,
								MountPath:     "/var/lib/tikv/raft",
							},
							{
								Name:          "newvol",
								Type:          "gp3",
								SizeGi:        100,
								Iops:          3000,
								ThroughputMbs: 125,
								MountPath:     "/var/lib/tikv/newvol",
							},
						},
					},
				},
			},
			wantErr:   true,
			assertion: func(c *cloudv1.Cluster) {},
		},
		"set tiproxy toggle workflow annotation and keep replicas unchanged when disabling tiproxy": {
			args: args{
				cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiProxy = cloudv1.TiProxy{
						Replicas:     1,
						InstanceType: "2c4g",
					}
				}),
				newConfigRevision: "",
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIPROXY,
						InstanceType:  "4c8g",
						Count:         0,
					},
					TidbGroupId: defaultLoadBalancerID,
				},
			},
			wantErr:            false,
			wantCompToggleType: componentToggleTypeDisable,
			assertion: func(c *cloudv1.Cluster) {
				assert.Equal(t, 1, c.Spec.Cluster.TiProxy.Replicas) // should not change
				assert.Contains(t, c.Annotations, consts.AnnoKeyTiProxyToggleWorkflow)
			},
		},
		"disable tiflash should directly set its spec": {
			args: args{
				cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiFlash = cloudv1.TiFlash{
						Replicas:     2,
						InstanceType: "2c4g",
					}
				}),
				newConfigRevision: "",
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIFLASH,
						InstanceType:  "2c4g",
						Count:         0,
					},
					TidbGroupId: defaultLoadBalancerID,
				},
			},
			wantErr:            false,
			wantCompToggleType: componentToggleTypeDisable,
			assertion: func(c *cloudv1.Cluster) {
				assert.Equal(t, 0, c.Spec.Cluster.TiFlash.Replicas)
				assert.NotContains(t, c.Annotations, consts.AnnoKeyTiProxyToggleWorkflow)
			},
		},
		"set config revision name if provided": {
			args: args{
				cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
					c.Spec.Cluster.TiDB = cloudv1.TiDB{
						Replicas:     1,
						InstanceType: "2c4g",
					}
					c.Spec.ConfigRevisionRef = &cloudv1.Reference{Name: "old-cr"}
				}),
				newConfigRevision: "new-cr",
				req: &apiv1.ScaleClusterRequest{
					ClusterId: existClusterId,
					NodeSpec: &apiv1.Cluster_NodeSpec{
						ComponentType: apicommon.ComponentType_TIDB,
						InstanceType:  "4c8g",
						Count:         2,
					},
					TidbGroupId: defaultLoadBalancerID,
				},
			},
			wantErr: false,
			assertion: func(c *cloudv1.Cluster) {
				assert.Equal(t, "new-cr", c.Spec.ConfigRevisionRef.Name)
				assert.Equal(t, "4c8g", c.Spec.Cluster.TiDB.InstanceType)
				assert.Equal(t, 2, c.Spec.Cluster.TiDB.Replicas)

			},
		},
	}
	for name, tt := range tests {
		t.Run(name, func(t *testing.T) {
			compToggleType, err := validateAndApplyClusterScaleChangesToCR(ctx, tt.args.cluster, tt.args.newConfigRevision, tt.args.req)
			if tt.wantErr {
				assert.Error(t, err)
				return
			}
			assert.NoError(t, err)
			tt.assertion(tt.args.cluster)
			// Check the returned value as well
			assert.Equal(t, tt.wantCompToggleType, compToggleType)
		})
	}
}

func validateAndApplyClusterScaleChangesToCR(ctx context.Context, cluster *cloudv1.Cluster, desiredConfigRevisionName string, req *apiv1.ScaleClusterRequest) (string, error) {
	validation, err := validateScaleClusterForApply(zap.NewNop(), cluster, req)
	if err != nil {
		return "", err
	}
	return applyClusterScaleChangesToCR(ctx, cluster, desiredConfigRevisionName, req, validation.storagePlan)
}

func TestApplyClusterScaleChangesRequiresGKEVolumeAttributesClass(t *testing.T) {
	const (
		currentStorageClass           = "tikv-hyperdisk-balanced-3000-140"
		desiredStorageClass           = "tikv-hyperdisk-balanced-6000-140"
		currentAdditionalStorageClass = "tikv-hyperdisk-balanced-3000-200"
		desiredAdditionalStorageClass = "tikv-hyperdisk-balanced-6000-200"
		pdSSDStorageClass             = "tikv-pd-ssd"
	)
	newPrimaryStorageRequest := func() *apiv1.ScaleClusterRequest {
		return &apiv1.ScaleClusterRequest{
			ClusterId: existClusterId,
			NodeSpec: &apiv1.Cluster_NodeSpec{
				ComponentType: apicommon.ComponentType_TIKV,
				Count:         3,
				Storage: &apiv1.Cluster_Storage{
					Type:          consts.GCPVolumeTypeHyperdiskBalanced,
					SizeGi:        100,
					Iops:          6000,
					ThroughputMbs: 140,
				},
			},
		}
	}
	newCluster := func(enableVAC bool) *cloudv1.Cluster {
		cluster := buildScaleTestGKECluster(func(cluster *cloudv1.Cluster) {
			cluster.Spec.Cluster.TiKV = cloudv1.TiKV{
				Replicas:         3,
				InstanceType:     "old-instance-type",
				StorageClassName: currentStorageClass,
				StorageSize:      resource.MustParse("100Gi"),
				StorageVolumes: []cloudv1.ComponentStorageVolume{
					{
						Name:             "raft",
						StorageClassName: currentAdditionalStorageClass,
						StorageSize:      resource.MustParse("100Gi"),
						MountPath:        "/var/lib/tikv/raft",
					},
				},
			}
			if enableVAC {
				cluster.EnableFeature(consts.ClusterFeatureUseVolumeAttributesClass)
			}
		})
		cluster.Status.DB.TiKV.Volume = &cloudv1.VolumeStatus{CurrentStorageClass: currentStorageClass}
		cluster.Status.DB.TiKV.StorageVolumes = []cloudv1.StorageVolumeStatus{{
			Name: "raft",
			VolumeStatus: cloudv1.VolumeStatus{
				CurrentStorageClass: currentAdditionalStorageClass,
			},
		}}
		return cluster
	}
	newAdditionalStorageRequest := func() *apiv1.ScaleClusterRequest {
		return &apiv1.ScaleClusterRequest{
			ClusterId: existClusterId,
			NodeSpec: &apiv1.Cluster_NodeSpec{
				ComponentType: apicommon.ComponentType_TIKV,
				InstanceType:  "new-instance-type",
				Count:         4,
				AdditionalStorage: []*apiv1.Cluster_AdditionalStorage{
					{
						Name:          "raft",
						Type:          consts.GCPVolumeTypeHyperdiskBalanced,
						SizeGi:        100,
						Iops:          6000,
						ThroughputMbs: 200,
						MountPath:     "/var/lib/tikv/raft",
					},
				},
			},
		}
	}

	t.Run("rejects before mutating Cluster spec", func(t *testing.T) {
		cluster := newCluster(false)
		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", newPrimaryStorageRequest())
		assert.Equal(t, codes.FailedPrecondition, status.Code(err))
		assert.Equal(t, currentStorageClass, cluster.Spec.Cluster.TiKV.StorageClassName)
		assert.Equal(t, "rev1", cluster.Spec.ConfigRevisionRef.Name)
	})

	t.Run("rejects additional volume before mutating Cluster spec", func(t *testing.T) {
		cluster := newCluster(false)
		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", newAdditionalStorageRequest())
		assert.Equal(t, codes.FailedPrecondition, status.Code(err))
		assert.Equal(t, 3, cluster.Spec.Cluster.TiKV.Replicas)
		assert.Equal(t, "old-instance-type", cluster.Spec.Cluster.TiKV.InstanceType)
		assert.Equal(t, currentAdditionalStorageClass, cluster.Spec.Cluster.TiKV.StorageVolumes[0].StorageClassName)
		assert.Equal(t, "rev1", cluster.Spec.ConfigRevisionRef.Name)
	})

	t.Run("rejects a primary StorageClass change when volume status is unavailable", func(t *testing.T) {
		tests := []struct {
			name        string
			clearStatus func(*cloudv1.Cluster)
		}{
			{name: "database status", clearStatus: func(cluster *cloudv1.Cluster) { cluster.Status.DB = nil }},
			{name: "TiKV status", clearStatus: func(cluster *cloudv1.Cluster) { cluster.Status.DB.TiKV = nil }},
			{name: "volume status", clearStatus: func(cluster *cloudv1.Cluster) { cluster.Status.DB.TiKV.Volume = nil }},
		}
		for _, tt := range tests {
			t.Run(tt.name, func(t *testing.T) {
				cluster := newCluster(true)
				tt.clearStatus(cluster)

				_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", newPrimaryStorageRequest())
				assert.Equal(t, codes.Unavailable, status.Code(err))
				assert.Equal(t, currentStorageClass, cluster.Spec.Cluster.TiKV.StorageClassName)
				assert.Equal(t, "rev1", cluster.Spec.ConfigRevisionRef.Name)
			})
		}
	})

	t.Run("allows a primary volume size change when volume status is unavailable", func(t *testing.T) {
		cluster := newCluster(true)
		cluster.Status.DB.TiKV.Volume = nil
		req := newPrimaryStorageRequest()
		req.NodeSpec.Storage = &apiv1.Cluster_Storage{
			Type:          consts.GCPVolumeTypeHyperdiskBalanced,
			SizeGi:        200,
			Iops:          3000,
			ThroughputMbs: 140,
		}

		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", req)
		assert.NoError(t, err)
		assert.Equal(t, currentStorageClass, cluster.Spec.Cluster.TiKV.StorageClassName)
		assert.Equal(t, "200Gi", cluster.Spec.Cluster.TiKV.StorageSize.String())
	})

	t.Run("allows TiFlash scale from zero without existing volume status", func(t *testing.T) {
		cluster := buildScaleTestGKECluster(func(cluster *cloudv1.Cluster) {
			cluster.EnableFeature(consts.ClusterFeatureUseVolumeAttributesClass)
			cluster.Spec.Cluster.TiFlash = cloudv1.TiFlash{Replicas: 0}
		})
		req := &apiv1.ScaleClusterRequest{
			ClusterId: existClusterId,
			NodeSpec: &apiv1.Cluster_NodeSpec{
				ComponentType: apicommon.ComponentType_TIFLASH,
				Count:         1,
				Storage: &apiv1.Cluster_Storage{
					Type:          consts.GCPVolumeTypeHyperdiskBalanced,
					SizeGi:        100,
					Iops:          6000,
					ThroughputMbs: 140,
				},
			},
		}

		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", req)
		assert.NoError(t, err)
		assert.Equal(t, 1, cluster.Spec.Cluster.TiFlash.Replicas)
		assert.Equal(t, "tiflash-hyperdisk-balanced-6000-140", cluster.Spec.Cluster.TiFlash.StorageClassName)
	})

	t.Run("ignores storage validation for a component without persistent storage", func(t *testing.T) {
		cluster := buildScaleTestGKECluster(func(cluster *cloudv1.Cluster) {
			cluster.EnableFeature(consts.ClusterFeatureUseVolumeAttributesClass)
			cluster.Spec.Cluster.TiDB = cloudv1.TiDB{Replicas: 1}
		})
		req := &apiv1.ScaleClusterRequest{
			ClusterId: existClusterId,
			NodeSpec: &apiv1.Cluster_NodeSpec{
				ComponentType: apicommon.ComponentType_TIDB,
				Count:         2,
				Storage: &apiv1.Cluster_Storage{
					Type:          consts.GCPVolumeTypeHyperdiskBalanced,
					SizeGi:        100,
					Iops:          6000,
					ThroughputMbs: 140,
				},
			},
		}

		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", req)
		assert.NoError(t, err)
		assert.Equal(t, 2, cluster.Spec.Cluster.TiDB.Replicas)
	})

	t.Run("allows supported modification", func(t *testing.T) {
		cluster := newCluster(true)
		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", newPrimaryStorageRequest())
		assert.NoError(t, err)
		assert.Equal(t, desiredStorageClass, cluster.Spec.Cluster.TiKV.StorageClassName)
		assert.Equal(t, "new-revision", cluster.Spec.ConfigRevisionRef.Name)
	})

	t.Run("rejects primary performance modification with dedicated raft-engine storage", func(t *testing.T) {
		cluster := newCluster(true)
		legacyStorageClass := pdSSDStorageClass
		cluster.Spec.Cluster.TiKV.DedicatedRaftEngineStorage = &cloudv1.TiKVDedicatedRaftEngineStorage{
			ComponentStorageSpec: cloudv1.ComponentStorageSpec{
				StorageSize:      resource.MustParse("100Gi"),
				StorageClassName: &legacyStorageClass,
			},
		}

		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", newPrimaryStorageRequest())
		assert.Equal(t, codes.FailedPrecondition, status.Code(err))
		assert.Equal(t, currentStorageClass, cluster.Spec.Cluster.TiKV.StorageClassName)
		assert.Equal(t, "rev1", cluster.Spec.ConfigRevisionRef.Name)
	})

	t.Run("allows primary size modification with dedicated raft-engine storage", func(t *testing.T) {
		cluster := newCluster(true)
		cluster.Spec.Cluster.TiKV.DedicatedRaftEngineStorage = &cloudv1.TiKVDedicatedRaftEngineStorage{}
		req := newPrimaryStorageRequest()
		req.NodeSpec.Storage = &apiv1.Cluster_Storage{
			Type:          consts.GCPVolumeTypeHyperdiskBalanced,
			SizeGi:        200,
			Iops:          3000,
			ThroughputMbs: 140,
		}

		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", req)
		assert.NoError(t, err)
		assert.Equal(t, currentStorageClass, cluster.Spec.Cluster.TiKV.StorageClassName)
		assert.Equal(t, "200Gi", cluster.Spec.Cluster.TiKV.StorageSize.String())
	})

	t.Run("allows an additional volume inheriting the desired primary type", func(t *testing.T) {
		cluster := newCluster(true)
		cluster.Spec.Cluster.TiKV.StorageVolumes[0].StorageClassName = ""

		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", newPrimaryStorageRequest())
		assert.NoError(t, err)
		assert.Equal(t, desiredStorageClass, cluster.Spec.Cluster.TiKV.StorageClassName)
	})

	t.Run("allows supported additional volume modification", func(t *testing.T) {
		cluster := newCluster(true)
		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", newAdditionalStorageRequest())
		assert.NoError(t, err)
		assert.Equal(t, desiredAdditionalStorageClass, cluster.Spec.Cluster.TiKV.StorageVolumes[0].StorageClassName)
		assert.Equal(t, "new-revision", cluster.Spec.ConfigRevisionRef.Name)
	})

	t.Run("rejects an additional StorageClass change when volume status is unavailable", func(t *testing.T) {
		cluster := newCluster(true)
		cluster.Status.DB.TiKV.StorageVolumes = nil

		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", newAdditionalStorageRequest())
		assert.Equal(t, codes.Unavailable, status.Code(err))
		assert.Equal(t, currentAdditionalStorageClass, cluster.Spec.Cluster.TiKV.StorageVolumes[0].StorageClassName)
		assert.Equal(t, "rev1", cluster.Spec.ConfigRevisionRef.Name)
	})

	t.Run("rejects an additional StorageClass change with an active VAC", func(t *testing.T) {
		cluster := newCluster(true)
		cluster.Status.DB.TiKV.StorageVolumes = []cloudv1.StorageVolumeStatus{{
			Name: "raft",
			VolumeStatus: cloudv1.VolumeStatus{
				CurrentStorageClass:          currentAdditionalStorageClass,
				CurrentVolumeAttributesClass: currentAdditionalStorageClass,
			},
		}}
		req := newAdditionalStorageRequest()
		req.NodeSpec.AdditionalStorage[0] = &apiv1.Cluster_AdditionalStorage{
			Name: "raft", Type: consts.GCPVolumeTypePDSSD, SizeGi: 100, MountPath: "/var/lib/tikv/raft",
		}

		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", req)
		assert.Equal(t, codes.FailedPrecondition, status.Code(err))
		assert.Equal(t, currentAdditionalStorageClass, cluster.Spec.Cluster.TiKV.StorageVolumes[0].StorageClassName)
		assert.Equal(t, "rev1", cluster.Spec.ConfigRevisionRef.Name)
	})

	t.Run("allows additional volume size change without VAC", func(t *testing.T) {
		cluster := newCluster(false)
		req := newAdditionalStorageRequest()
		req.NodeSpec.AdditionalStorage[0].Iops = 3000
		req.NodeSpec.AdditionalStorage[0].SizeGi = 200
		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", req)
		assert.NoError(t, err)
		assert.Equal(t, currentAdditionalStorageClass, cluster.Spec.Cluster.TiKV.StorageVolumes[0].StorageClassName)
		assert.Equal(t, "200Gi", cluster.Spec.Cluster.TiKV.StorageVolumes[0].StorageSize.String())
	})

	t.Run("rejects operator v2", func(t *testing.T) {
		cluster := newCluster(true)
		cluster.EnableFeature(consts.ClusterFeatureUseTiDBOperatorV2)
		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", newPrimaryStorageRequest())
		assert.Equal(t, codes.FailedPrecondition, status.Code(err))
		assert.Equal(t, currentStorageClass, cluster.Spec.Cluster.TiKV.StorageClassName)
		assert.Equal(t, "rev1", cluster.Spec.ConfigRevisionRef.Name)
	})

	t.Run("allows Hyperdisk to pd-ssd without an active VAC", func(t *testing.T) {
		cluster := newCluster(true)
		req := newPrimaryStorageRequest()
		req.NodeSpec.Storage = &apiv1.Cluster_Storage{Type: consts.GCPVolumeTypePDSSD, SizeGi: 100}

		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", req)
		assert.NoError(t, err)
		assert.Equal(t, pdSSDStorageClass, cluster.Spec.Cluster.TiKV.StorageClassName)
		assert.Equal(t, "new-revision", cluster.Spec.ConfigRevisionRef.Name)
	})

	t.Run("rejects Hyperdisk to pd-ssd with an active VAC", func(t *testing.T) {
		cluster := newCluster(true)
		cluster.Status.DB = &cloudv1.DBStatus{TiKV: &cloudv1.TiKVStatus{Volume: &cloudv1.VolumeStatus{
			CurrentStorageClass:          currentStorageClass,
			CurrentVolumeAttributesClass: currentStorageClass,
		}}}
		req := newPrimaryStorageRequest()
		req.NodeSpec.Storage = &apiv1.Cluster_Storage{Type: consts.GCPVolumeTypePDSSD, SizeGi: 100}

		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", req)
		assert.Equal(t, codes.FailedPrecondition, status.Code(err))
		assert.Equal(t, currentStorageClass, cluster.Spec.Cluster.TiKV.StorageClassName)
		assert.Equal(t, "rev1", cluster.Spec.ConfigRevisionRef.Name)
	})

	t.Run("rejects a mixed topology when primary Hyperdisk performance changes", func(t *testing.T) {
		cluster := newCluster(true)
		cluster.Spec.Cluster.TiKV.StorageVolumes[0].StorageClassName = pdSSDStorageClass

		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", newPrimaryStorageRequest())
		assert.Equal(t, codes.InvalidArgument, status.Code(err))
		assert.Equal(t, currentStorageClass, cluster.Spec.Cluster.TiKV.StorageClassName)
		assert.Equal(t, "rev1", cluster.Spec.ConfigRevisionRef.Name)
	})

	t.Run("rejects migrating an additional volume to Hyperdisk while primary performance changes", func(t *testing.T) {
		cluster := newCluster(true)
		cluster.Spec.Cluster.TiKV.StorageVolumes[0].StorageClassName = pdSSDStorageClass
		req := newPrimaryStorageRequest()
		req.NodeSpec.AdditionalStorage = []*apiv1.Cluster_AdditionalStorage{
			{
				Name:          "raft",
				Type:          consts.GCPVolumeTypeHyperdiskBalanced,
				SizeGi:        100,
				Iops:          6000,
				ThroughputMbs: 200,
				MountPath:     "/var/lib/tikv/raft",
			},
		}

		_, err := validateAndApplyClusterScaleChangesToCR(context.Background(), cluster, "new-revision", req)
		assert.Equal(t, codes.InvalidArgument, status.Code(err))
		assert.Equal(t, currentStorageClass, cluster.Spec.Cluster.TiKV.StorageClassName)
		assert.Equal(t, pdSSDStorageClass, cluster.Spec.Cluster.TiKV.StorageVolumes[0].StorageClassName)
		assert.Equal(t, "rev1", cluster.Spec.ConfigRevisionRef.Name)
	})
}

func TestScaleClusterRejectsInvalidGKEStorageBeforeConfigUpdate(t *testing.T) {
	ctx := context.Background()
	cluster := buildScaleTestGKECluster(func(cluster *cloudv1.Cluster) {
		cluster.Spec.Cluster.TiKV = cloudv1.TiKV{
			Replicas:         3,
			InstanceType:     "2c4g",
			StorageClassName: "tikv-hyperdisk-balanced-3000-140",
			StorageSize:      resource.MustParse("100Gi"),
		}
	})
	existingConfigRevision := &cloudv1.ConfigRevision{
		ObjectMeta: metav1.ObjectMeta{Name: cluster.Spec.ConfigRevisionRef.Name},
		Spec:       cloudv1.ConfigRevisionSpec{Version: "v7.5.1"},
	}
	cli := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithObjects(
		cluster, existingConfigRevision, buildInstanceTypeConfigMap(),
	).Build()
	server := &ClusterServer{
		instanceTypeProvider: NewInstanceTypeProvider(),
		KubeClient:           cli,
		clusterReader:        &util.LabelSelectorReader{},
	}

	_, err := server.ScaleCluster(ctx, &apiv1.ScaleClusterRequest{
		ClusterId: existClusterId,
		NodeSpec: &apiv1.Cluster_NodeSpec{
			ComponentType: apicommon.ComponentType_TIKV,
			Count:         3,
			Storage: &apiv1.Cluster_Storage{
				Type:          consts.GCPVolumeTypeHyperdiskBalanced,
				SizeGi:        100,
				Iops:          6000,
				ThroughputMbs: 140,
			},
		},
		ComponentConfigurations: []*apiv1.ComponentConfiguration{{
			ComponentType:             apicommon.ComponentType_TIKV,
			ConfigurationFileContents: `foo = "bar"`,
		}},
	})
	assert.Equal(t, codes.FailedPrecondition, status.Code(err))

	var configRevisions cloudv1.ConfigRevisionList
	assert.NoError(t, cli.List(ctx, &configRevisions))
	assert.Len(t, configRevisions.Items, 1)
	assert.Equal(t, existingConfigRevision.Name, configRevisions.Items[0].Name)
}

func TestSetTiproxyToggleWfIntoAnnoIfNeed(t *testing.T) {
	tests := []struct {
		name              string
		cluster           *cloudv1.Cluster
		tiproxyToggleType string
		groupId           string
		wantErr           bool
		clusterAssertions func(*cloudv1.Cluster)
	}{
		{
			name: "no toggle type should return nil",
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Annotations = map[string]string{}
			}),
			tiproxyToggleType: componentToggleTypeNone,
			groupId:           defaultLoadBalancerID,
			wantErr:           false,
			clusterAssertions: func(c *cloudv1.Cluster) {
				assert.NotContains(t, c.Annotations, consts.AnnoKeyTiProxyToggleWorkflow)
			},
		},
		{
			name: "enable tiproxy should set workflow annotation",
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Annotations = map[string]string{}
			}),
			tiproxyToggleType: componentToggleTypeEnable,
			groupId:           defaultLoadBalancerID,
			wantErr:           false,
			clusterAssertions: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, consts.AnnoKeyTiProxyToggleWorkflow)
			},
		},
		{
			name: "disable tiproxy should set workflow annotation",
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Annotations = map[string]string{}
			}),
			tiproxyToggleType: componentToggleTypeDisable,
			groupId:           defaultLoadBalancerID,
			wantErr:           false,
			clusterAssertions: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, consts.AnnoKeyTiProxyToggleWorkflow)
			},
		},
		{
			name: "workflow already running should return nil",
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Annotations = map[string]string{
					consts.AnnoKeyTiProxyToggleWorkflow: "metadata:\n  name: test-workflow",
				}
			}),
			tiproxyToggleType: componentToggleTypeEnable,
			groupId:           defaultLoadBalancerID,
			wantErr:           false,
			clusterAssertions: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations[consts.AnnoKeyTiProxyToggleWorkflow], "name: test-workflow")
			},
		},
		{
			name: "invalid workflow annotation should return error",
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Annotations = map[string]string{
					consts.AnnoKeyTiProxyToggleWorkflow: "invalid yaml",
				}
			}),
			tiproxyToggleType: componentToggleTypeEnable,
			groupId:           defaultLoadBalancerID,
			wantErr:           true,
			clusterAssertions: func(c *cloudv1.Cluster) {
				assert.Equal(t, c.Annotations[consts.AnnoKeyTiProxyToggleWorkflow], "invalid yaml")
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := setTiproxyToggleWfIntoAnnoIfNeed(zap.NewNop(), tt.cluster, tt.tiproxyToggleType, tt.groupId, cloudv1.ComponentTiProxy)
			if tt.wantErr {
				assert.Error(t, err)
				return
			}
			assert.NoError(t, err)
			tt.clusterAssertions(tt.cluster)
		})
	}
}

func TestSetScaleClusterWfIntoAnnoIfNeed(t *testing.T) {
	ctx := context.Background()
	mockProvider := cloudprovider.NewInstanceTypeProviderFromList(ctx, []cloudv1.InstanceType{
		{
			Name:   "",
			CPU:    resource.MustParse("0"),
			Memory: resource.MustParse("0"),
		},
		{
			Name:   "2c4g",
			CPU:    resource.MustParse("2"),
			Memory: resource.MustParse("4Gi"),
		},
		{
			Name:   "4c8g",
			CPU:    resource.MustParse("4"),
			Memory: resource.MustParse("8Gi"),
		},
	})

	tests := map[string]struct {
		cluster      *cloudv1.Cluster
		changedComp  cloudv1.ClusterComponent
		newConfigRev string
		setup        func(*cloudv1.Cluster)
		wantErr      bool
		assertions   func(*cloudv1.Cluster)
	}{
		"should set workflow annotation if not present": {
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Annotations = map[string]string{}
				c.Spec.Cluster.TiDB = cloudv1.TiDB{Replicas: 2, InstanceType: "2c4g"}
			}),
			changedComp: &cloudv1.TiDB{
				Replicas:     3,
				InstanceType: "4c8g",
			},
			newConfigRev: "my-config-rev",
			wantErr:      false,
			assertions: func(c *cloudv1.Cluster) {
				assert.Contains(t, c.Annotations, consts.AnnoKeyScaleClusterWorkflow)
				val := c.Annotations[consts.AnnoKeyScaleClusterWorkflow]
				assert.Contains(t, val, "my-config-rev")
			},
		},
		"should not overwrite existing workflow annotation": {
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Annotations = map[string]string{
					consts.AnnoKeyScaleClusterWorkflow: "metadata:\n  name: test-existing-workflow",
				}
			}),
			changedComp: &cloudv1.TiDB{
				Replicas:     3,
				InstanceType: "4c8g",
			},
			newConfigRev: "new-cr",
			wantErr:      false,
			assertions: func(c *cloudv1.Cluster) {
				assert.Equal(t, "metadata:\n  name: test-existing-workflow", c.Annotations[consts.AnnoKeyScaleClusterWorkflow])
			},
		},
		"should return error if annotation is invalid yaml": {
			cluster: buildTestClusterCR(existClusterId, func(c *cloudv1.Cluster) {
				c.Annotations = map[string]string{
					consts.AnnoKeyScaleClusterWorkflow: "invalid yaml",
				}
			}),
			changedComp: &cloudv1.TiDB{
				Replicas:     3,
				InstanceType: "4c8g",
			},
			newConfigRev: "new-cr",
			wantErr:      true,
			assertions: func(c *cloudv1.Cluster) {
				assert.Equal(t, "invalid yaml", c.Annotations[consts.AnnoKeyScaleClusterWorkflow])
			},
		},
	}

	for name, tt := range tests {
		t.Run(name, func(t *testing.T) {
			if tt.setup != nil {
				tt.setup(tt.cluster)
			}
			componentWrapper := &cloudv1.ClusterComponentSpecWrapper{
				ClusterComponent: tt.changedComp,
			}
			err := setScaleClusterWfIntoAnnoIfNeed(context.Background(), zap.NewNop(), mockProvider, tt.cluster, componentWrapper, tt.newConfigRev, false)
			if tt.wantErr {
				assert.Error(t, err)
				tt.assertions(tt.cluster)
				return
			}
			assert.NoError(t, err)
			tt.assertions(tt.cluster)
		})
	}
}

func Test_getComponentToggleType(t *testing.T) {
	tests := []struct {
		name string
		comp cloudv1.ClusterComponent
		req  *apiv1.ScaleClusterRequest
		want string
	}{
		{
			name: "enable tiproxy from 0 to 1",
			comp: &cloudv1.TiProxy{Replicas: 0},
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIPROXY,
					Count:         1,
				},
			},
			want: componentToggleTypeEnable,
		},
		{
			name: "disable tiproxy from 1 to 0",
			comp: &cloudv1.TiProxy{Replicas: 1},
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIPROXY,
					Count:         0,
				},
			},
			want: componentToggleTypeDisable,
		},
		{
			name: "disable tiflash from 1 to 0",
			comp: &cloudv1.TiFlash{Replicas: 1},
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIFLASH,
					Count:         0,
				},
			},
			want: componentToggleTypeDisable,
		},
		{
			name: "scale tiproxy from 1 to 2 should be none",
			comp: &cloudv1.TiProxy{Replicas: 1},
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIPROXY,
					Count:         2,
				},
			},
			want: componentToggleTypeNone,
		},
		{
			name: "scale tiproxy from 2 to 1 should be none",
			comp: &cloudv1.TiProxy{Replicas: 2},
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIPROXY,
					Count:         1,
				},
			},
			want: componentToggleTypeNone,
		},
		{
			name: "scale tidb from 0 to 1 should be Enabled", // for example, create tidb in extra group
			comp: &cloudv1.TiDB{Replicas: 0},
			req: &apiv1.ScaleClusterRequest{
				NodeSpec: &apiv1.Cluster_NodeSpec{
					ComponentType: apicommon.ComponentType_TIDB,
					Count:         1,
				},
			},
			want: componentToggleTypeEnable,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewWithT(t)
			got := getComponentToggleType(tt.comp, tt.req)
			g.Expect(got).To(Equal(tt.want))
		})
	}
}

func NewInstanceTypeProvider() gitops.InstanceTypeProvider {
	ctx := context.Background()
	return cloudprovider.NewInstanceTypeProviderFromList(ctx, []cloudv1.InstanceType{
		{
			Name:   "",
			CPU:    resource.MustParse("0"),
			Memory: resource.MustParse("0"),
		},
		{
			Name:   "2c4g",
			CPU:    resource.MustParse("2"),
			Memory: resource.MustParse("4Gi"),
		},
		{
			Name:   "4c8g",
			CPU:    resource.MustParse("4"),
			Memory: resource.MustParse("8Gi"),
		},
		{
			Name:   "8c16g",
			CPU:    resource.MustParse("4"),
			Memory: resource.MustParse("8Gi"),
		},
	})
}