package tidbautoscaling

import (
	"context"
	"fmt"
	"regexp"
	"testing"
	"time"

	"github.com/go-logr/logr"
	. "github.com/onsi/gomega"
	serverlessv1alpha1 "github.com/tidbcloud/aws-shared-provider/apis/serverless/v1alpha1"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/clientfake"
	infra "github.com/tidbcloud/tidb-management-service/api/serverless_infra"
	"github.com/tidbcloud/tidb-management-service/api/serverless_infra/meta"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

func TestTidbAutoscaling_scale(t *testing.T) {
	type tokenReplicas struct {
		token    string
		replicas int
	}
	type tokenScaleRecord struct {
		token string
		scaleRecord
	}
	type tokenWantScaleRecord struct {
		token           string
		direction       int
		updateTimestamp bool
	}

	tests := []struct {
		name             string
		cluster          *infra.Cluster
		tokenMetricMap   map[string]tokenMetric
		taConfig         *serverlessv1alpha1.TidbAutoscaling
		tidbMaxReplicas  int
		scaleRecord      *tokenScaleRecord
		disableScaleDown map[string]bool
		wantScaleRecord  *tokenWantScaleRecord
		wantScale        *tokenReplicas
	}{
		{
			name:    "not validateScale, replicas == 0",
			cluster: &infra.Cluster{ObjectMeta: meta.ObjectMeta{Name: "token1"}},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.4, true}, {10, true}},
				}, 0},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
		},
		{
			name:    "not validateScale, metric 2 not validate",
			cluster: &infra.Cluster{ObjectMeta: meta.ObjectMeta{Name: "token1"}},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.4, true}, {10, false}},
				}, 1},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
		},
		{
			name: "scale up to max replicas with scaleRecord",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.4, true}, {10, true}},
				}, 1},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
			scaleRecord: &tokenScaleRecord{
				token:       "token1",
				scaleRecord: scaleRecord{HOLD, time.Now().Add(-1 * time.Second).Unix()},
			},
			wantScaleRecord: &tokenWantScaleRecord{"token1", UP, true},
			wantScale:       &tokenReplicas{"token1", 10},
		},
		{
			name: "scale up to Infra MaxReplicas with scaleRecord",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Spec:       infra.ClusterSpec{TiDB: infra.TiDBSpec{MaxReplicas: ptr.To(10)}},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.4, true}, {10, true}},
				}, 1},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 2,
			scaleRecord: &tokenScaleRecord{
				token:       "token1",
				scaleRecord: scaleRecord{HOLD, time.Now().Add(-1 * time.Second).Unix()},
			},
			wantScaleRecord: &tokenWantScaleRecord{"token1", UP, true},
			wantScale:       &tokenReplicas{"token1", 10},
		},
		{
			name: "only metric 2 scale up",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.4, true}, {2, true}},
					"pod2": []float64OrNil{{0.6, true}, {4, true}},
				}, 2},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
			wantScaleRecord: &tokenWantScaleRecord{"token1", UP, true},
			wantScale:       &tokenReplicas{"token1", 6},
		},
		{
			name: "all scale up, use metric 1",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.8, true}, {3, true}},
					"pod2": []float64OrNil{{1.2, true}, {5, true}},
				}, 2},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.5, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1.1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
			wantScaleRecord: &tokenWantScaleRecord{"token1", UP, true},
			wantScale:       &tokenReplicas{"token1", 8},
		},
		{
			name: "resource group scale up",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Spec:       infra.ClusterSpec{TiDB: infra.TiDBSpec{ResourceGroups: map[string]*infra.ResourceGroup{"rg1": {}}}},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1/rg1": {podMetrics{
					"pod1": []float64OrNil{{2, true}},
					"pod2": []float64OrNil{{4, true}},
				}, 2},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				Metrics: []serverlessv1alpha1.Metric{{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"}},
			},
			tidbMaxReplicas: 10,
			wantScaleRecord: &tokenWantScaleRecord{"token1/rg1", UP, true},
			wantScale:       &tokenReplicas{"token1/rg1", 6},
		},
		{
			name:    "no scale for flapping",
			cluster: &infra.Cluster{ObjectMeta: meta.ObjectMeta{Name: "token1"}},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.7, true}, {1, true}},
					"pod2": []float64OrNil{{0.6, true}, {1.1, true}},
				}, 2},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
		},
		{
			name: "scale down, use metric 2",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.1, true}, {0.3, true}},
					"pod2": []float64OrNil{{0.1, true}, {0.4, true}},
					"pod3": []float64OrNil{{0.1, true}, {0.5, true}},
				}, 3},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
			wantScaleRecord: &tokenWantScaleRecord{"token1", DOWN, true},
			wantScale:       &tokenReplicas{"token1", 2},
		},
		{
			name: "scale down for resource group",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Spec:       infra.ClusterSpec{TiDB: infra.TiDBSpec{ResourceGroups: map[string]*infra.ResourceGroup{"rg1": {}}}},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1/rg1": {podMetrics{
					"pod1": []float64OrNil{{0.3, true}},
					"pod2": []float64OrNil{{0.4, true}},
					"pod3": []float64OrNil{{0.5, true}},
				}, 3},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				Metrics: []serverlessv1alpha1.Metric{{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"}},
			},
			tidbMaxReplicas: 10,
			wantScaleRecord: &tokenWantScaleRecord{"token1/rg1", DOWN, true},
			wantScale:       &tokenReplicas{"token1/rg1", 2},
		},
		{
			name: "scale down, use metric 1 with scaleRecord",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.1, true}, {0.3, true}},
					"pod2": []float64OrNil{{0.1, true}, {0.4, true}},
					"pod3": []float64OrNil{{0.1, true}, {0.5, true}},
				}, 3},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
			scaleRecord: &tokenScaleRecord{
				token:       "token1",
				scaleRecord: scaleRecord{UP, time.Now().Add(-1 * time.Second).Unix()},
			},
			wantScaleRecord: &tokenWantScaleRecord{"token1", DOWN, true},
			wantScale:       &tokenReplicas{"token1", 2},
		},
		{
			name: "no scale up for stabilization window, create scaleRecord",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{3, true}, {2, true}},
				}, 1},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				ScaleUpStabilizationWindowSeconds: 10,
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
			wantScaleRecord: &tokenWantScaleRecord{"token1", UP, true},
		},
		{
			name: "no scale up for stabilization window, update scaleRecord",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{3, true}, {2, true}},
				}, 1},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				ScaleUpStabilizationWindowSeconds: 10,
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
			scaleRecord: &tokenScaleRecord{
				token:       "token1",
				scaleRecord: scaleRecord{HOLD, time.Now().Add(-1 * time.Second).Unix()},
			},
			wantScaleRecord: &tokenWantScaleRecord{"token1", UP, true},
		},
		{
			name:    "no scale up for stabilization window, not update scaleRecord",
			cluster: &infra.Cluster{ObjectMeta: meta.ObjectMeta{Name: "token1"}},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{3, true}, {2, true}},
				}, 1},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				ScaleUpStabilizationWindowSeconds: 10,
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
			scaleRecord: &tokenScaleRecord{
				token:       "token1",
				scaleRecord: scaleRecord{UP, time.Now().Unix()},
			},
			wantScaleRecord: &tokenWantScaleRecord{"token1", UP, false},
		},
		{
			name: "scale up for stabilization window",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{3, true}, {2, true}},
				}, 1},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				ScaleUpStabilizationWindowSeconds: 10,
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
			scaleRecord: &tokenScaleRecord{
				token:       "token1",
				scaleRecord: scaleRecord{UP, time.Now().Add(-10 * time.Second).Unix()},
			},
			wantScaleRecord: &tokenWantScaleRecord{"token1", UP, true},
			wantScale:       &tokenReplicas{"token1", 5},
		},
		{
			name: "no scale down because of Infra MinReplicas",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Spec:       infra.ClusterSpec{TiDB: infra.TiDBSpec{MinReplicas: ptr.To(3)}},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.1, true}, {0.3, true}},
					"pod2": []float64OrNil{{0.1, true}, {0.4, true}},
					"pod3": []float64OrNil{{0.1, true}, {0.5, true}},
				}, 3},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
		},
		{
			name: "no scale down because of disableScaleDown",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.1, true}, {0.3, true}},
					"pod2": []float64OrNil{{0.1, true}, {0.4, true}},
					"pod3": []float64OrNil{{0.1, true}, {0.5, true}},
				}, 3},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 0.7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 1, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas:  10,
			disableScaleDown: map[string]bool{"token1": true},
		},
		{
			name: "no scale down for stabilization window, create scaleRecord",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.1, true}, {0.2, true}},
					"pod2": []float64OrNil{{0.2, true}, {0.1, true}},
				}, 2},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				ScaleDownStabilizationWindowSeconds: 10,
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 10, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
			wantScaleRecord: &tokenWantScaleRecord{"token1", DOWN, true},
		},
		{
			name: "no scale down for stabilization window, update scaleRecord",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.1, true}, {0.2, true}},
					"pod2": []float64OrNil{{0.2, true}, {0.1, true}},
				}, 2},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				ScaleDownStabilizationWindowSeconds: 10,
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 10, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
			scaleRecord: &tokenScaleRecord{
				token:       "token1",
				scaleRecord: scaleRecord{UP, time.Now().Add(-1 * time.Second).Unix()},
			},
			wantScaleRecord: &tokenWantScaleRecord{"token1", DOWN, true},
		},
		{
			name:    "no scale down for stabilization window, not update scaleRecord",
			cluster: &infra.Cluster{ObjectMeta: meta.ObjectMeta{Name: "token1"}},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.1, true}, {0.2, true}},
					"pod2": []float64OrNil{{0.2, true}, {0.1, true}},
				}, 2},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				ScaleDownStabilizationWindowSeconds: 10,
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 10, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
			scaleRecord: &tokenScaleRecord{
				token:       "token1",
				scaleRecord: scaleRecord{DOWN, time.Now().Unix()},
			},
			wantScaleRecord: &tokenWantScaleRecord{"token1", DOWN, false},
		},
		{
			name: "scale down for stabilization window",
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "token1"},
				Status:     infra.ClusterStatus{KeyspaceToken: "token1"},
			},
			tokenMetricMap: map[string]tokenMetric{
				"token1": {podMetrics{
					"pod1": []float64OrNil{{0.1, true}, {0.2, true}},
					"pod2": []float64OrNil{{0.2, true}, {0.1, true}},
				}, 2},
			},
			taConfig: &serverlessv1alpha1.TidbAutoscaling{
				ScaleDownStabilizationWindowSeconds: 10,
				Metrics: []serverlessv1alpha1.Metric{
					{AvgValue: 7, ScaleFlapping: 0.1, PodLabel: "pod"},
					{AvgValue: 10, ScaleFlapping: 0.1, PodLabel: "pod"},
				},
			},
			tidbMaxReplicas: 10,
			scaleRecord: &tokenScaleRecord{
				token:       "token1",
				scaleRecord: scaleRecord{DOWN, time.Now().Add(-10 * time.Second).Unix()},
			},
			wantScaleRecord: &tokenWantScaleRecord{"token1", DOWN, true},
			wantScale:       &tokenReplicas{"token1", 1},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			ta := &TidbAutoscaling{
				Logger:         logr.New(log.NullLogSink{}),
				ctx:            context.Background(),
				scaleRecordMap: map[string]scaleRecord{},
				infraClusterCli: clientfake.NewInfraClient(tt.cluster).
					WithPatch(func(set map[string]any, obj infra.Object) error {
						if obj.(*infra.Cluster).Annotations == nil {
							obj.(*infra.Cluster).Annotations = map[string]string{}
						}
						path := fmt.Sprintf("annotations.'%s'", types.AnnoKeyTidbAutoscaling)
						obj.(*infra.Cluster).Annotations[types.AnnoKeyTidbAutoscaling] = set[path].(string)

						re := regexp.MustCompile(`^'spec\.tidb'\.resource_groups\.'([^']+)'\.replicas$`)
						for k, v := range set {
							if k == "'spec.tidb'.replicas" {
								obj.(*infra.Cluster).Spec.TiDB.Replicas = v.(int)
							}
							if m := re.FindStringSubmatch(k); len(m) == 2 {
								obj.(*infra.Cluster).Spec.TiDB.ResourceGroups[m[1]].Replicas = v.(int)
							}
						}
						return nil
					}).
					WithIndex("status.keyspace_token-index", func(obj infra.Object) (value, dst any) {
						cluster := obj.(*infra.Cluster)
						return cluster.Status.KeyspaceToken, meta.ObjectKey{Namespace: cluster.Namespace, Name: cluster.Name}
					}).
					Build(),
			}
			if tt.scaleRecord != nil {
				ta.scaleRecordMap[tt.scaleRecord.token] = tt.scaleRecord.scaleRecord
			}
			ta.scale(tt.tokenMetricMap, tt.taConfig, tt.tidbMaxReplicas, tt.disableScaleDown)

			if tt.wantScaleRecord == nil {
				g.Expect(ta.scaleRecordMap).To(BeEmpty())
			} else {
				g.Expect(ta.scaleRecordMap).To(HaveKey(tt.wantScaleRecord.token))
				g.Expect(ta.scaleRecordMap[tt.wantScaleRecord.token].direction).To(Equal(tt.wantScaleRecord.direction))
				if tt.scaleRecord != nil {
					if tt.wantScaleRecord.updateTimestamp {
						g.Expect(ta.scaleRecordMap[tt.wantScaleRecord.token].timestamp).To(BeNumerically(">", tt.scaleRecord.timestamp))
					} else {
						g.Expect(ta.scaleRecordMap[tt.wantScaleRecord.token].timestamp).To(Equal(tt.scaleRecord.timestamp))
					}
				}
			}

			cluster := &infra.Cluster{ObjectMeta: meta.ObjectMeta{Namespace: tt.cluster.Namespace, Name: tt.cluster.Name}}
			g.Expect(ta.infraClusterCli.Get(ta.ctx, cluster)).To(Succeed())
			if tt.wantScale == nil {
				g.Expect(cluster.Annotations).NotTo(HaveKey(types.AnnoKeyTidbAutoscaling))
			} else {
				g.Expect(cluster.Annotations).To(HaveKey(types.AnnoKeyTidbAutoscaling))
				for token := range tt.tokenMetricMap {
					if rgName := utils.TrimResourceGroup(token); rgName != "" {
						g.Expect(cluster.Spec.TiDB.ResourceGroups).To(HaveKey(rgName))
						g.Expect(cluster.Spec.TiDB.ResourceGroups[rgName].Replicas).To(Equal(tt.wantScale.replicas))
					} else {
						g.Expect(cluster.Spec.TiDB.Replicas).To(Equal(tt.wantScale.replicas))
					}
				}
			}
		})
	}
}