package cluster

import (
	"strconv"
	"strings"
	"testing"
	"time"

	. "github.com/onsi/gomega"
	"github.com/pingcap/errors"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbmanager"
	"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"
	ep "github.com/tidbcloud/aws-shared-provider/pkg/utils/endpoint"
	infra "github.com/tidbcloud/tidb-management-service/api/serverless_infra"
	"github.com/tidbcloud/tidb-management-service/api/serverless_infra/meta"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
)

func TestGenPublicEndpointStatus(t *testing.T) {
	t.Setenv("NAMESPACE", "tidb-admin")

	cases := map[string]struct {
		cluster        *infra.Cluster
		observed       []runtime.Object
		expectEndpoint *infra.PublicStatus
	}{
		"CreatePublicEndpoint": {
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{
					Name:      "Test-CreatePublicEndpoint",
					Namespace: "test-ns",
				},
				Spec: infra.ClusterSpec{
					SqlEndpoints: infra.SqlEndpointsSpec{
						Public: &infra.PublicSpec{Enabled: true},
					},
				},
			},
			observed: []runtime.Object{
				&corev1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb-gateway-new",
						Namespace: "tidb-admin",
						Annotations: map[string]string{
							types.AnnoGatewayDnsName: "gateway01.us-east-1.dev.shared.aws.tidbcloud.com",
						},
					},
				},
			},
			expectEndpoint: &infra.PublicStatus{
				Host: "gateway01.us-east-1.dev.shared.aws.tidbcloud.com",
				Port: defaultTiDBPort,
			},
		},
	}

	httpServer := newFakeHTTP()
	defer httpServer.Close()

	for name, c := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)

			r := newFakeReconcileContext(c.cluster, newFakeReconciler(c.observed, []infra.Object{c.cluster}, httpServer), nil)
			publicEndpoint, err := r.genPublicEndpointStatus()
			g.Expect(err).To(Succeed())
			g.Expect(publicEndpoint).To(Equal(c.expectEndpoint))
		})
	}
}

func TestGenPrivatelinkStatus(t *testing.T) {
	t.Setenv("NAMESPACE", "tidb-admin")

	cases := map[string]struct {
		cluster           *infra.Cluster
		observed          []runtime.Object
		expectPrivatelink *infra.AwsPrivatelinkServiceStatus
	}{
		"CreatePrivatelink": {
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{
					Name:      "Test-CreatePrivatelink",
					Namespace: "test-ns",
				},
				Spec: infra.ClusterSpec{
					SqlEndpoints: infra.SqlEndpointsSpec{
						AwsPrivatelinkService: &infra.AwsPrivatelinkServiceSpec{Enabled: true},
					},
				},
			},
			observed: []runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "privatelink-service-config",
						Namespace: "tidb-admin",
					},
					Data: map[string]string{
						"private-dns":        "gateway01-privatelink.us-east-1.dev.shared.aws.tidbcloud.com",
						"service-name":       "com.amazonaws.vpce.us-east-1.vpce-svc-0699cfb9b1db28a22",
						"availability-zones": "us-east-1a,us-west-2b",
					},
				},
			},
			expectPrivatelink: &infra.AwsPrivatelinkServiceStatus{
				AvailabilityZones: []string{"us-east-1a", "us-west-2b"},
				ServiceName:       "com.amazonaws.vpce.us-east-1.vpce-svc-0699cfb9b1db28a22",
				PrivateDnsName:    "gateway01-privatelink.us-east-1.dev.shared.aws.tidbcloud.com",
				Port:              defaultTiDBPort,
			},
		},
		"CreatePrivatelinkZR": {
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{
					Name:      "Test-CreatePrivatelink-ZR",
					Namespace: "test-ns",
				},
				Spec: infra.ClusterSpec{
					SqlEndpoints: infra.SqlEndpointsSpec{
						AwsPrivatelinkService: &infra.AwsPrivatelinkServiceSpec{Enabled: true},
					},
					EnableZoneRedundant: "yes",
				},
			},
			observed: []runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "privatelink-zr-service-config",
						Namespace: "tidb-admin",
					},
					Data: map[string]string{
						"private-dns":        "gateway01-privatelink-zr.us-east-1.dev.shared.aws.tidbcloud.com",
						"service-name":       "com.amazonaws.vpce.us-east-1.vpce-svc-9999cfb9c1db28a88",
						"availability-zones": "us-east-1c,us-west-2b",
					},
				},
			},
			expectPrivatelink: &infra.AwsPrivatelinkServiceStatus{
				AvailabilityZones: []string{"us-east-1c", "us-west-2b"},
				ServiceName:       "com.amazonaws.vpce.us-east-1.vpce-svc-9999cfb9c1db28a88",
				PrivateDnsName:    "gateway01-privatelink-zr.us-east-1.dev.shared.aws.tidbcloud.com",
				Port:              defaultTiDBPort,
			},
		},
	}

	httpServer := newFakeHTTP()
	defer httpServer.Close()

	for name, c := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)

			r := newFakeReconcileContext(c.cluster, newFakeReconciler(c.observed, []infra.Object{c.cluster}, httpServer), nil)
			privatelink, err := r.genPrivatelinkStatus()
			g.Expect(err).To(Succeed())
			g.Expect(privatelink).To(Equal(c.expectPrivatelink))
		})
	}
}

func TestReconcileTidbPool(t *testing.T) {
	cases := map[string]struct {
		cluster               *infra.Cluster
		tidbManagers          map[string]tidbmanager.TidbManager
		etcdTokenPrefixMap    map[string]string
		expectAnno            bool
		expectOldPoolReplicas map[string]int
		expectErr             error
	}{
		"PoolUnchanged": {
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "test", Namespace: "ns"},
				Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{
					Replicas: 1, PoolSelector: defaultTidbPool,
					ResourceGroups: map[string]*infra.ResourceGroup{"rg1": {Replicas: 2}},
				}},
				Status: infra.ClusterStatus{
					Phase:          infra.ClusterPhaseRunning,
					TidbPool:       defaultTidbPool,
					TiDB:           infra.TiDBStatus{UpdatedReplicas: 1},
					ResourceGroups: map[string]infra.ResourceGroup{"rg1": {Replicas: 2}},
				},
			},
			etcdTokenPrefixMap: map[string]string{
				defaultTidbPool: "/default-pool/",
			},
		},
		"PoolStartChanging": {
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "test", Namespace: "ns"},
				Spec:       infra.ClusterSpec{TiDB: infra.TiDBSpec{Replicas: 2, PoolSelector: "new-pool"}},
				Status: infra.ClusterStatus{
					Phase:         infra.ClusterPhaseRunning,
					TidbPool:      defaultTidbPool,
					TiDB:          infra.TiDBStatus{UpdatedReplicas: 2},
					KeyspaceToken: "keyspace-token",
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod1,ip1"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod2,ip2"}},
					},
				},
				"new-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod3,ip3"}},
					},
				},
			},
			etcdTokenPrefixMap: map[string]string{
				defaultTidbPool: "/default-pool/",
				"other-pool":    "/other-pool/",
				"new-pool":      "/new-pool/",
			},
			expectAnno:            true,
			expectOldPoolReplicas: map[string]int{"keyspace-token": 2},
			expectErr:             &utils.RequeueAfterError{},
		},
		"PoolStartChanging_ResourceGroup": {
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "test", Namespace: "ns"},
				Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{
					Replicas:       1,
					ResourceGroups: map[string]*infra.ResourceGroup{"rg1": {Pool: "new-pool", Replicas: 2}},
				}},
				Status: infra.ClusterStatus{
					Phase:          infra.ClusterPhaseRunning,
					TidbPool:       defaultTidbPool,
					TiDB:           infra.TiDBStatus{UpdatedReplicas: 1},
					KeyspaceToken:  "keyspace-token",
					ResourceGroups: map[string]infra.ResourceGroup{"rg1": {Pool: defaultTidbPool, Replicas: 2}},
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod1,ip1"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod2,ip2"}},
					},
				},
				"new-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod3,ip3"}},
					},
				},
			},
			etcdTokenPrefixMap: map[string]string{
				defaultTidbPool: "/default-pool/",
				"other-pool":    "/other-pool/",
				"new-pool":      "/new-pool/",
			},
			expectAnno:            true,
			expectOldPoolReplicas: map[string]int{"keyspace-token/rg1": 2},
			expectErr:             &utils.RequeueAfterError{},
		},
		"PoolStartChanging_Both": {
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{Name: "test", Namespace: "ns"},
				Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{
					Replicas:       1,
					PoolSelector:   "new-pool",
					ResourceGroups: map[string]*infra.ResourceGroup{"rg1": {Replicas: 2}},
				}},
				Status: infra.ClusterStatus{
					Phase:          infra.ClusterPhaseRunning,
					TidbPool:       defaultTidbPool,
					TiDB:           infra.TiDBStatus{UpdatedReplicas: 1},
					KeyspaceToken:  "keyspace-token",
					ResourceGroups: map[string]infra.ResourceGroup{"rg1": {Pool: defaultTidbPool, Replicas: 2}},
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token":     {Endpoints: []string{"tidb-server-pod0,ip0"}},
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod1,ip1"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod2,ip2"}},
					},
				},
				"new-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod3,ip3"}},
					},
				},
			},
			etcdTokenPrefixMap: map[string]string{
				defaultTidbPool: "/default-pool/",
				"other-pool":    "/other-pool/",
				"new-pool":      "/new-pool/",
			},
			expectAnno:            true,
			expectOldPoolReplicas: map[string]int{"keyspace-token/rg1": 2, "keyspace-token": 1},
			expectErr:             &utils.RequeueAfterError{},
		},
		"PoolChanging_WaitNewTidbPoolRunning": {
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{
					Name:      "test",
					Namespace: "ns",
					Annotations: map[string]string{
						types.AnnoKeyTidbPoolChanging: time.Now().Add(-retryTidbPoolChanging).Format(time.RFC3339),
					},
				},
				Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{Replicas: 2, PoolSelector: "new-pool"}},
				Status: infra.ClusterStatus{
					Phase:         infra.ClusterPhaseRunning,
					TidbPool:      "new-pool",
					TiDB:          infra.TiDBStatus{UpdatedReplicas: 2},
					KeyspaceToken: "keyspace-token",
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod1,ip1"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod2,ip2"}},
					},
				},
			},
			etcdTokenPrefixMap: map[string]string{
				defaultTidbPool: "/default-pool/",
				"other-pool":    "/other-pool/",
				"new-pool":      "/new-pool/",
			},
			expectAnno:            true,
			expectOldPoolReplicas: map[string]int{"keyspace-token": 2},
			expectErr: &utils.RequeueAfterError{
				Cause:        errors.New("waiting new tidb pool pods to running"),
				RequeueAfter: waitNewTidbPoolRunning,
			},
		},
		"PoolChanging_WaitNewTidbPoolRunning_ResourceGroup": {
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{
					Name:      "test",
					Namespace: "ns",
					Annotations: map[string]string{
						types.AnnoKeyTidbPoolChanging: time.Now().Add(-retryTidbPoolChanging).Format(time.RFC3339),
					},
				},
				Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{
					Replicas:       1,
					PoolSelector:   "new-pool",
					ResourceGroups: map[string]*infra.ResourceGroup{"rg1": {Replicas: 2}},
				}},
				Status: infra.ClusterStatus{
					Phase:          infra.ClusterPhaseRunning,
					TidbPool:       "new-pool",
					TiDB:           infra.TiDBStatus{UpdatedReplicas: 1},
					KeyspaceToken:  "keyspace-token",
					ResourceGroups: map[string]infra.ResourceGroup{"rg1": {Pool: "new-pool", Replicas: 2}},
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod1,ip1"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod2,ip2"}},
					},
				},
			},
			etcdTokenPrefixMap: map[string]string{
				defaultTidbPool: "/default-pool/",
				"other-pool":    "/other-pool/",
				"new-pool":      "/new-pool/",
			},
			expectAnno:            true,
			expectOldPoolReplicas: map[string]int{"keyspace-token/rg1": 2},
			expectErr: &utils.RequeueAfterError{
				Cause:        errors.New("waiting new tidb pool pods to running"),
				RequeueAfter: waitNewTidbPoolRunning,
			},
		},
		"PoolChanging_ReduceOne": {
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{
					Name:      "test",
					Namespace: "ns",
					Annotations: map[string]string{
						types.AnnoKeyTidbPoolChanging: time.Now().Add(-retryTidbPoolChanging).Format(time.RFC3339),
					},
				},
				Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{Replicas: 3, PoolSelector: "new-pool"}},
				Status: infra.ClusterStatus{
					Phase:         infra.ClusterPhaseRunning,
					TidbPool:      "new-pool",
					TiDB:          infra.TiDBStatus{UpdatedReplicas: 3},
					KeyspaceToken: "keyspace-token",
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod1,ip1", "tidb-server-pod4,ip4"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod2,ip2"}},
					},
				},
				"new-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod3,ip3"}},
					},
				},
			},
			etcdTokenPrefixMap: map[string]string{
				defaultTidbPool: "/default-pool/",
				"other-pool":    "/other-pool/",
				"new-pool":      "/new-pool/",
			},
			expectAnno:            true,
			expectOldPoolReplicas: map[string]int{"keyspace-token": 2},
			expectErr:             &utils.RequeueAfterError{RequeueAfter: retryTidbPoolChanging},
		},
		"PoolChanging_ReduceOne_ResourceGroup": {
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{
					Name:      "test",
					Namespace: "ns",
					Annotations: map[string]string{
						types.AnnoKeyTidbPoolChanging: time.Now().Add(-retryTidbPoolChanging).Format(time.RFC3339),
					},
				},
				Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{
					Replicas:       1,
					ResourceGroups: map[string]*infra.ResourceGroup{"rg1": {Pool: "new-pool", Replicas: 3}},
				}},
				Status: infra.ClusterStatus{
					Phase:          infra.ClusterPhaseRunning,
					TidbPool:       defaultTidbPool,
					TiDB:           infra.TiDBStatus{UpdatedReplicas: 1},
					KeyspaceToken:  "keyspace-token",
					ResourceGroups: map[string]infra.ResourceGroup{"rg1": {Pool: "new-pool", Replicas: 3}},
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod1,ip1", "tidb-server-pod4,ip4"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod2,ip2"}},
					},
				},
				"new-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod3,ip3"}},
					},
				},
			},
			etcdTokenPrefixMap: map[string]string{
				defaultTidbPool: "/default-pool/",
				"other-pool":    "/other-pool/",
				"new-pool":      "/new-pool/",
			},
			expectAnno:            true,
			expectOldPoolReplicas: map[string]int{"keyspace-token/rg1": 2},
			expectErr:             &utils.RequeueAfterError{RequeueAfter: retryTidbPoolChanging},
		},
		"PoolEndChanging": {
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{
					Name:      "test",
					Namespace: "ns",
					Annotations: map[string]string{
						types.AnnoKeyTidbPoolChanging: time.Now().Format(time.RFC3339),
					},
				},
				Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{
					Replicas: 1, PoolSelector: defaultTidbPool,
					ResourceGroups: map[string]*infra.ResourceGroup{"rg1": {Pool: "other-pool", Replicas: 1}},
				}},
				Status: infra.ClusterStatus{
					Phase:          infra.ClusterPhaseRunning,
					TidbPool:       defaultTidbPool,
					TiDB:           infra.TiDBStatus{UpdatedReplicas: 1},
					KeyspaceToken:  "keyspace-token",
					ResourceGroups: map[string]infra.ResourceGroup{"rg1": {Pool: "other-pool", Replicas: 1}},
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod1,ip1"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token":     {},
						"keyspace-token/rg1": {Endpoints: []string{"tidb-server-pod2,ip2"}},
					},
				},
			},
			etcdTokenPrefixMap: map[string]string{
				defaultTidbPool: "/default-pool/",
				"other-pool":    "/other-pool/",
			},
		},
	}

	httpServer := newFakeHTTP()
	defer httpServer.Close()

	for name, c := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)

			tidbPods := make([]runtime.Object, 0)
			etcdDB := make(map[string]string)
			for pool, tManager := range c.tidbManagers {
				for keyspaceToken, eps := range tManager.(*clientfake.FakeTidbManager).GetAssignedTidbByTokenR {
					etcdDB[c.etcdTokenPrefixMap[pool]+keyspaceToken+"/"+utils.KeyReplicas] = strconv.Itoa(len(eps.Endpoints))
					etcdDB[c.etcdTokenPrefixMap[pool]+keyspaceToken+"/"+utils.KeyEndpoints] = strings.Join(eps.Endpoints, ";")
					for _, endpoint := range eps.Endpoints {
						podName, podIP := ep.ParsePodFromEndpoint(endpoint)
						tidbPod := &corev1.Pod{
							ObjectMeta: metav1.ObjectMeta{
								Name:        podName,
								Namespace:   pool,
								Annotations: make(map[string]string),
							},
							Status: corev1.PodStatus{Phase: corev1.PodRunning, PodIP: podIP},
						}
						// intentionally setting up AnnoTidbStatusRunSlow in the new pool
						// if pool == c.cluster.Spec.TiDB.PoolSelector {
						// 	tidbPod.Annotations[types.AnnoTidbStatusRunSlow] = "true"
						// }
						tidbPods = append(tidbPods, tidbPod)
					}
				}
			}

			nr := newFakeReconciler(tidbPods, []infra.Object{c.cluster}, httpServer)
			nr.TidbManagers = c.tidbManagers
			nr.EtcdKV = clientfake.NewEtcdKV(clientfake.BuildArg{DB: etcdDB})
			nr.EtcdTokenPrefixMap = c.etcdTokenPrefixMap
			nr.TidbPools = []string{defaultTidbPool}
			for pool := range c.etcdTokenPrefixMap {
				if pool != defaultTidbPool {
					nr.TidbPools = append(nr.TidbPools, pool)
				}
			}
			r := newFakeReconcileContext(c.cluster, nr, nil)

			err := r.reconcileTidbPool()
			if c.expectErr != nil {
				if reqErr, ok := c.expectErr.(*utils.RequeueAfterError); ok && reqErr.RequeueAfter == 0 {
					g.Expect(err.(*utils.RequeueAfterError).RequeueAfter).To(BeNumerically("<", retryTidbPoolChanging))
				} else {
					g.Expect(err).To(MatchError(c.expectErr.Error()))
				}
			} else {
				g.Expect(err).To(Succeed())
			}
			g.Expect(c.cluster.Status.TidbPool).To(Equal(r.specTidbPool))
			g.Expect(r.cluster.Status.TiDB.UpdatedReplicas).To(Equal(r.cluster.Spec.TiDB.Replicas))
			for rgName, rgSpec := range c.cluster.Spec.TiDB.ResourceGroups {
				rgSpecCopy := rgSpec.DeepCopy()
				if rgSpecCopy.Pool == "" {
					rgSpecCopy.Pool = r.specTidbPool
				}
				g.Expect(r.cluster.Status.ResourceGroups[rgName]).To(Equal(*rgSpecCopy))
			}

			if c.expectAnno {
				g.Expect(r.cluster.Annotations).To(HaveKey(types.AnnoKeyTidbPoolChanging))

				rgNames := []string{""}
				for rgName := range r.cluster.Status.ResourceGroups {
					rgNames = append(rgNames, rgName)
				}
				for _, rgName := range rgNames {
					// verify the all old pools replicas
					oldPoolReplicas := 0
					token, pool := r.keyspaceToken, r.specTidbPool
					if rgName != "" {
						token, pool = r.keyspaceToken+"/"+rgName, r.cluster.Status.ResourceGroups[rgName].Pool
					}
					for _, etcdTokenPrefix := range r.EtcdTokenPrefixMap {
						if etcdTokenPrefix == r.EtcdTokenPrefixMap[pool] {
							continue
						}
						resp, err := r.EtcdKV.Get(r.ctx, etcdTokenPrefix+token+"/"+utils.KeyReplicas)
						g.Expect(err).To(Succeed())
						if resp.Count != 0 {
							r, err := strconv.Atoi(string(resp.Kvs[0].Value))
							g.Expect(err).To(Succeed())
							oldPoolReplicas += r
						}
					}
					g.Expect(oldPoolReplicas).To(Equal(c.expectOldPoolReplicas[token]), "token: %s", token)

					// verify the new pool replicas
					originalPool := r.specTidbPool
					if rgName != "" {
						if originalPool = r.cluster.Spec.TiDB.ResourceGroups[rgName].Pool; originalPool == "" {
							originalPool = r.specTidbPool
						}
					}
					if originalPool != pool {
						resp, err := r.EtcdKV.Get(r.ctx, r.EtcdTokenPrefixMap[pool]+token+"/"+utils.KeyReplicas)
						g.Expect(err).To(Succeed())
						g.Expect(string(resp.Kvs[0].Value)).To(Equal(strconv.Itoa(c.cluster.Spec.TiDB.Replicas)))

						// verify the old tidb pools pod annotations, should have AnnoTidbStatusRunSlow
						// for _, objKey := range r.allOldTidbPoolPods() {
						// 	tidbPod := &corev1.Pod{}
						// 	g.Expect(r.Client.Get(r.ctx, objKey, tidbPod)).To(Succeed())
						// 	g.Expect(tidbPod.Annotations).To(HaveKey(types.AnnoTidbStatusRunSlow))
						// }

						// verify the new tidb pool pod annotations, should not have AnnoTidbStatusRunSlow
						// for _, objKey := range r.newTidbPoolPods() {
						// 	tidbPod := &corev1.Pod{}
						// 	g.Expect(r.Client.Get(r.ctx, objKey, tidbPod)).To(Succeed())
						// 	g.Expect(tidbPod.Annotations).ToNot(HaveKey(types.AnnoTidbStatusRunSlow))
						// }
					}
				}
			} else {
				g.Expect(r.cluster.Annotations).ToNot(HaveKey(types.AnnoKeyTidbPoolChanging))
				// verify all old tidb pool paths are removed
				for _, etcdTokenPrefix := range r.EtcdTokenPrefixMap {
					if etcdTokenPrefix == r.EtcdTokenPrefixMap[r.specTidbPool] {
						continue
					}
					resp, err := r.EtcdKV.Get(r.ctx, etcdTokenPrefix+r.keyspaceToken+"/"+utils.KeyReplicas)
					g.Expect(err).To(Succeed())
					g.Expect(resp.Count).To(Equal(int64(0)))
				}
			}
		})
	}
}