//go:build unit

package eks

import (
	"context"
	"testing"
	"time"

	"github.com/crossplane-contrib/provider-aws/apis/autoscaling/v1beta1"
	ec2v1alpha1 "github.com/crossplane-contrib/provider-aws/apis/ec2/v1alpha1"
	commonv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	. "github.com/onsi/gomega"
	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	helmaction "helm.sh/helm/v3/pkg/action"
	apps "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/record"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	. "github.com/tidbcloud/eks-provider/controllers/upgradeoperation/common"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/dataplane"
	"github.com/tidbcloud/infra-provider/pkg/extratidb"
	"github.com/tidbcloud/infra-provider/pkg/gitops"
	"github.com/tidbcloud/infra-provider/pkg/namer"
	"github.com/tidbcloud/infra-provider/pkg/nodegroup/selfmanaged"
)

// fakeEnvProvider implements dataplane.EnvProvider using a fake k8s client.
type fakeEnvProvider struct {
	eksCli client.Client
}

func (f *fakeEnvProvider) GetClient() client.Client                 { return f.eksCli }
func (f *fakeEnvProvider) GetHelmConfig() *helmaction.Configuration { return nil }
func (f *fakeEnvProvider) GetRestConfig() *rest.Config              { return nil }
func (f *fakeEnvProvider) GetRegion() string                        { return "us-west-2" }
func (f *fakeEnvProvider) GetClusterName() string                   { return "test-eks-cluster" }
func (f *fakeEnvProvider) GetClusterEndpoint() string               { return "" }
func (f *fakeEnvProvider) GetClusterCA() []byte                     { return nil }
func (f *fakeEnvProvider) GetAccount() string                       { return "" }
func (f *fakeEnvProvider) GetNetwork() string                       { return "" }
func (f *fakeEnvProvider) AWS() dataplane.AWSInterface              { return nil }
func (f *fakeEnvProvider) GCP() dataplane.GCPInterface              { return nil }

var _ dataplane.EnvProvider = &fakeEnvProvider{}

// makeReadyASG creates a new-version ASG that passes allNodeGroupsReady and
// GetTiDBClusterSelfManagedNodeGroups (which filters by component label and instance type).
// The ASG references a launch template named "<name>-lt" with instance type "m5.xlarge".
// Use makeReadyASGLaunchTemplate to create the matching launch template in the infra client.
func makeReadyASG(name, clusterLabel, desiredVersion string) *v1beta1.AutoScalingGroup {
	ltName := name + "-lt"
	return &v1beta1.AutoScalingGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name: name,
			Labels: map[string]string{
				namer.ClusterSelectorKey:   clusterLabel,
				namer.ComponentSelectorKey: string(cloudv1.ComponentTiDB),
			},
			Annotations: map[string]string{
				selfmanaged.K8sVersionAnnotationKey: desiredVersion,
			},
		},
		Spec: v1beta1.AutoScalingGroupSpec{
			ForProvider: v1beta1.AutoScalingGroupParameters{
				DesiredCapacity: pointer.Int64(1),
				LaunchTemplate: &v1beta1.LaunchTemplateSpecification{
					LaunchTemplateName: &ltName,
				},
			},
		},
		Status: v1beta1.AutoScalingGroupStatus{
			ResourceStatus: commonv1.ResourceStatus{
				ConditionedStatus: commonv1.ConditionedStatus{
					Conditions: []commonv1.Condition{
						{
							Type:   commonv1.TypeReady,
							Status: corev1.ConditionTrue,
						},
						{
							Type:   commonv1.TypeSynced,
							Status: corev1.ConditionTrue,
						},
					},
				},
			},
		},
	}
}

// makeReadyASGLaunchTemplate creates the LaunchTemplate CR that makeReadyASG references.
// The instance type is "m5.xlarge" to match makeTestCluster's TiDB instance type.
func makeReadyASGLaunchTemplate(asgName string) *ec2v1alpha1.LaunchTemplate {
	instanceType := "m5.xlarge"
	ltName := asgName + "-lt"
	return &ec2v1alpha1.LaunchTemplate{
		ObjectMeta: metav1.ObjectMeta{
			Name: ltName,
		},
		Spec: ec2v1alpha1.LaunchTemplateSpec{
			ForProvider: ec2v1alpha1.LaunchTemplateParameters{
				LaunchTemplateData: &ec2v1alpha1.RequestLaunchTemplateData{
					InstanceType: &instanceType,
				},
			},
		},
	}
}

// makeTCWithRestartAnnotation creates a TidbCluster with the restart annotation already set
// so that MutateTCForRestarting returns false (mutation already applied).
// If ready=true, TidbClusterReady condition is True; otherwise False.
func makeTCWithRestartAnnotation(name, namespace, version string, ready bool) *opv1alpha1.TidbCluster {
	restartAnnotation := map[string]string{
		RestartTiDBAnnotation: version,
	}
	tc := &opv1alpha1.TidbCluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
		},
		Spec: opv1alpha1.TidbClusterSpec{
			PD: &opv1alpha1.PDSpec{
				ComponentSpec: opv1alpha1.ComponentSpec{
					Annotations: restartAnnotation,
				},
			},
			TiKV: &opv1alpha1.TiKVSpec{
				ComponentSpec: opv1alpha1.ComponentSpec{
					Annotations: restartAnnotation,
				},
			},
			TiDB: &opv1alpha1.TiDBSpec{
				ComponentSpec: opv1alpha1.ComponentSpec{
					Annotations: restartAnnotation,
				},
			},
		},
	}

	condStatus := corev1.ConditionFalse
	if ready {
		condStatus = corev1.ConditionTrue
	}
	tc.Status.Conditions = []opv1alpha1.TidbClusterCondition{
		{
			Type:   opv1alpha1.TidbClusterReady,
			Status: condStatus,
		},
	}
	if ready {
		tc.Spec.TiDB.Replicas = 1
		tc.Status.TiDB.Phase = opv1alpha1.NormalPhase
		tc.Status.TiDB.Members = map[string]opv1alpha1.TiDBMember{
			name + "-tidb-0": {Name: name + "-tidb-0", Health: true},
		}
		tc.Status.TiDB.StatefulSet = &apps.StatefulSetStatus{ReadyReplicas: 1}
	}

	return tc
}

func restartAnnot(version string) map[string]string {
	return map[string]string{RestartTiDBAnnotation: version}
}

// makeTCWithTiProxyPhaseComplete returns a TC that can finish the TiProxy restart phase in one reconcile.
func makeTCWithTiProxyPhaseComplete(name, namespace, version string) *opv1alpha1.TidbCluster {
	return &opv1alpha1.TidbCluster{
		ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: namespace},
		Spec: opv1alpha1.TidbClusterSpec{
			PD:   &opv1alpha1.PDSpec{Replicas: 0},
			TiKV: &opv1alpha1.TiKVSpec{Replicas: 0},
			TiDB: &opv1alpha1.TiDBSpec{Replicas: 0},
			TiProxy: &opv1alpha1.TiProxySpec{
				Replicas: 1,
				ComponentSpec: opv1alpha1.ComponentSpec{
					Annotations: restartAnnot(version),
				},
			},
		},
		Status: opv1alpha1.TidbClusterStatus{
			TiProxy: opv1alpha1.TiProxyStatus{
				Synced: true,
				Members: map[string]opv1alpha1.TiProxyMember{
					"tiproxy-0": {Health: true},
				},
			},
		},
	}
}

func podsOnASGNode(namespace, asgName, nodeSuffix string, pods map[string]string) ([]client.Object, *corev1.Node) {
	nodeName := "node-" + asgName + nodeSuffix
	node := getTestNode(nodeName, asgName, false)
	objs := []client.Object{node}
	for name, component := range pods {
		objs = append(objs, getTestPod(name, namespace, component, nodeName))
	}
	return objs, node
}

func makeTCWithTiProxyNoRestart(name, namespace string, ready bool) *opv1alpha1.TidbCluster {
	tc := &opv1alpha1.TidbCluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
		},
		Spec: opv1alpha1.TidbClusterSpec{
			PD: &opv1alpha1.PDSpec{
				ComponentSpec: opv1alpha1.ComponentSpec{},
			},
			TiKV: &opv1alpha1.TiKVSpec{
				ComponentSpec: opv1alpha1.ComponentSpec{},
			},
			TiDB: &opv1alpha1.TiDBSpec{
				ComponentSpec: opv1alpha1.ComponentSpec{},
			},
			TiProxy: &opv1alpha1.TiProxySpec{
				ComponentSpec: opv1alpha1.ComponentSpec{},
				Replicas:      1,
			},
		},
	}

	condStatus := corev1.ConditionFalse
	if ready {
		condStatus = corev1.ConditionTrue
	}
	tc.Status.Conditions = []opv1alpha1.TidbClusterCondition{
		{
			Type:   opv1alpha1.TidbClusterReady,
			Status: condStatus,
		},
	}

	return tc
}

// makeTCWithTiProxyAllRestartAnnotation creates a TidbCluster with restart annotations already set
// for all components including TiProxy.
func makeTCWithTiProxyAllRestartAnnotation(name, namespace, version string, ready bool) *opv1alpha1.TidbCluster {
	restartAnnotation := map[string]string{
		RestartTiDBAnnotation: version,
	}
	tc := &opv1alpha1.TidbCluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
		},
		Spec: opv1alpha1.TidbClusterSpec{
			PD: &opv1alpha1.PDSpec{
				ComponentSpec: opv1alpha1.ComponentSpec{
					Annotations: restartAnnotation,
				},
			},
			TiKV: &opv1alpha1.TiKVSpec{
				ComponentSpec: opv1alpha1.ComponentSpec{
					Annotations: restartAnnotation,
				},
			},
			TiDB: &opv1alpha1.TiDBSpec{
				ComponentSpec: opv1alpha1.ComponentSpec{
					Annotations: restartAnnotation,
				},
			},
			TiProxy: &opv1alpha1.TiProxySpec{
				ComponentSpec: opv1alpha1.ComponentSpec{
					Annotations: restartAnnotation,
				},
				Replicas: 1,
			},
		},
		Status: opv1alpha1.TidbClusterStatus{
			TiProxy: opv1alpha1.TiProxyStatus{
				Synced: true,
				Members: map[string]opv1alpha1.TiProxyMember{
					"tiproxy-0": {
						Health: true,
					},
				},
			},
		},
	}

	condStatus := corev1.ConditionFalse
	if ready {
		condStatus = corev1.ConditionTrue
	}
	tc.Status.Conditions = []opv1alpha1.TidbClusterCondition{
		{
			Type:   opv1alpha1.TidbClusterReady,
			Status: condStatus,
		},
	}

	return tc
}

// makeTestCluster creates a minimal Cluster CR for testing.
// The cluster has no EKS status (FromProvider.Raw == nil), so GetTiDBNodeVersion returns ""
// which is different from desiredVersion, making allTiDBNodesUpgraded return false.
// TiDB instance type is set to "m5.xlarge" to match makeReadyASG.
func makeTestCluster(name, desiredVersion string, extraGroupName string) *cloudv1.Cluster {
	cluster := &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name: name,
			Labels: map[string]string{
				cloudv1.ProviderTypeLabel: consts.ProviderTypeEKS,
			},
		},
		Spec: cloudv1.ClusterSpec{
			Cluster: cloudv1.ClusterTopology{
				TiDB: cloudv1.TiDB{
					Replicas:     1,
					InstanceType: "m5.xlarge",
				},
			},
		},
	}

	if extraGroupName != "" {
		cluster.Spec.ExtraGroups = []cloudv1.ExtraGroupSpec{
			{
				Name: extraGroupName,
				TiDB: cloudv1.TiDB{
					Replicas: 1,
				},
			},
		}
	}

	return cluster
}

// makeTestUpgradeOperation creates a minimal UpgradeOperation for testing.
func makeTestUpgradeOperation(name, namespace, desiredVersion string) *cloudv1.UpgradeOperation {
	return &cloudv1.UpgradeOperation{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
		},
		Status: cloudv1.UpgradeOperationStatus{
			CurrentStep: &cloudv1.Step{
				Version: desiredVersion,
				Options: map[string]string{
					// Disable scale-out extra tidb to avoid the extra tidb scale-in path.
					ScaleOutExtraTiDBAnnotation: "false",
				},
			},
		},
	}
}

// Test_UpgradeTiDBNodes_ExtraGroupTCReadiness_BugCondition is the canonical bug condition exploration
// test for Bug 1, named as required by the spec task.
//
// Bug 1: In UpgradeTiDBNodes, after all restart mutations are applied and all pods are on new nodes
// and the default TC is ready, the function proceeds to scale-in WITHOUT checking if extra group TCs
// are ready.
//
// Setup (Bug Condition C1):
//   - All restart annotations already set → MutateTCForRestartingWithoutTiProxy and MutateTCForRestarting return false
//   - All TiDB pods on new ASG nodes → areAllTiDBPodsOnNewNodeGroups returns true
//   - Default TC: TidbClusterReady=True → IsTCReady(&tc) passes
//   - Extra group TC: TidbClusterReady=False → should trigger requeue on fixed code
//
// Expected (correct) behavior: return ctrl.Result{Requeue: true} — wait for extra group TC.
// Actual (unfixed) behavior: proceed to UpdateClusterStatus, return ctrl.Result{RequeueAfter: 1s}.
//
// This test FAILS on unfixed code (Requeue=false), confirming the bug exists.
// It PASSES on fixed code (Requeue=true), confirming the fix works.
//
// Counterexample documented: "UpgradeTiDBNodes proceeds to UpdateClusterStatus (RequeueAfter=1s)
// even though extra group TC is not ready (TidbClusterReady=False). result.Requeue is false."
//
// **Validates: Requirements 2.1, 2.2**
func Test_UpgradeTiDBNodes_ExtraGroupTCReadiness_BugCondition(t *testing.T) {
	const (
		desiredVersion = "1.29"
		clusterName    = "bugtest-cluster"
		extraGroupName = "analytics"
	)
	namespace := cloudv1.ClusterPrefixId(clusterName)

	ctx := context.Background()
	g := NewWithT(t)

	// Build the cluster CR with one non-deleted extra group.
	cluster := makeTestCluster(clusterName, desiredVersion, extraGroupName)

	// Compute the extra group TC name using the same namer logic as upgrader.go.
	extraGroupNamer := namer.NewNamer().
		ForTiDBGroup(cluster, &cluster.Spec.ExtraGroups[0], "").
		WithComponent(string(cloudv1.ComponentTiDB))
	extraGroupTCName := extraGroupNamer.GetTiDBName()

	// UpgradeOperation with scaleOutExtraTiDB=false to avoid the extra tidb scale-in path.
	uo := makeTestUpgradeOperation("bug-uo", namespace, desiredVersion)

	// --- Infra client: one new-version ASG (numNewNGs >= numOldNGs, allNodeGroupsReady=true) ---
	newASG := makeReadyASG("tidb-ng-new-bug", clusterName, desiredVersion)

	infraScheme := runtime.NewScheme()
	g.Expect(v1beta1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(cloudv1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(ec2v1alpha1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())

	infraCli := fake.NewClientBuilder().
		WithScheme(infraScheme).
		WithObjects(newASG, makeReadyASGLaunchTemplate(newASG.Name), cluster, uo).
		WithStatusSubresource(cluster).
		WithStatusSubresource(uo).
		Build()

	// --- EKS client: TCs, pods, nodes ---
	// Default TC: restart annotation set at desiredVersion → MutateTCForRestartingWithoutTiProxy returns false.
	//             TidbClusterReady=True → IsTCReady(&tc) passes.
	defaultTC := makeTCWithRestartAnnotation(consts.ClusterPrefix, namespace, desiredVersion, true)

	// Extra group TC: restart annotation set at desiredVersion → MutateTCForRestarting returns false.
	//                 TidbClusterReady=False → this is the bug condition: unfixed code never checks this.
	extraGroupTC := makeTCWithRestartAnnotation(extraGroupTCName, namespace, desiredVersion, false)

	// One TiDB pod on the new ASG node → areAllTiDBPodsOnNewNodeGroups returns true.
	tidbPod := getTestPod("db-tidb-0", namespace, string(cloudv1.ComponentTiDB), "node-tidb-bug")
	tidbNode := getTestNode("node-tidb-bug", newASG.Name, false) // self-managed: "Name" label

	eksScheme := runtime.NewScheme()
	g.Expect(corev1.AddToScheme(eksScheme)).Should(Succeed())
	g.Expect(opv1alpha1.AddToScheme(eksScheme)).Should(Succeed())

	eksCli := fake.NewClientBuilder().
		WithScheme(eksScheme).
		WithObjects(defaultTC, extraGroupTC, tidbPod, tidbNode).
		Build()

	u := &upgrader{
		infraCli:      infraCli,
		eventRecorder: record.NewFakeRecorder(100),
	}
	envProvider := &fakeEnvProvider{eksCli: eksCli}

	// --- Call UpgradeTiDBNodes ---
	result, err := u.UpgradeTiDBNodes(ctx, uo, envProvider, cluster)

	// Fixed code: waits for extra group TC readiness before proceeding (Phase 7).
	g.Expect(err).Should(BeNil(), "expected no error from UpgradeTiDBNodes")
	g.Expect(result.Requeue).Should(BeFalse(),
		"UpgradeTiDBNodes should not proceed while extra group TC %q is not ready; got result=%+v",
		extraGroupTCName, result)
	g.Expect(result.RequeueAfter).Should(Equal(2*time.Minute),
		"expected RequeueAfter=2m while waiting for extra group TC readiness, got %v",
		result.RequeueAfter)
}

// Test_UpgradeTiDBNodes_Preservation_NoExtraGroups verifies that when the cluster has no extra
// groups, UpgradeTiDBNodes behaves identically to the original: the new extra-group readiness
// loop is a no-op and the function proceeds past the readiness gate to UpdateClusterStatus,
// returning ctrl.Result{RequeueAfter: 1s}.
//
// This test PASSES on UNFIXED code — it establishes the preservation baseline for Bug 1.
//
// Setup: all restart mutations applied, all pods on new nodes, default TC ready, no extra groups.
// Expected: function proceeds past readiness gate → returns RequeueAfter=1s (NOT Requeue=true).
//
// **Validates: Requirements 3.1, 3.2, 3.3, 3.4**
func Test_UpgradeTiDBNodes_Preservation_NoExtraGroups(t *testing.T) {
	const (
		desiredVersion = "1.29"
		clusterName    = "preserve-no-extra"
	)
	namespace := cloudv1.ClusterPrefixId(clusterName)

	ctx := context.Background()
	g := NewWithT(t)

	// Cluster with NO extra groups.
	cluster := makeTestCluster(clusterName, desiredVersion, "")

	// UpgradeOperation with scaleOutExtraTiDB=false to skip the extra tidb scale-in path.
	uo := makeTestUpgradeOperation("preserve-uo-no-extra", namespace, desiredVersion)

	// One new-version ASG: numNewNGs(1) >= numOldNGs(0) → skip creation; allNodeGroupsReady → true.
	newASG := makeReadyASG("tidb-ng-preserve-no-extra", clusterName, desiredVersion)

	infraScheme := runtime.NewScheme()
	g.Expect(v1beta1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(cloudv1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(ec2v1alpha1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())

	infraCli := fake.NewClientBuilder().
		WithScheme(infraScheme).
		WithObjects(newASG, makeReadyASGLaunchTemplate(newASG.Name), cluster, uo).
		WithStatusSubresource(cluster).
		WithStatusSubresource(uo).
		Build()

	// Default TC: restart annotation already set → MutateTCForRestartingWithoutTiProxy returns false.
	//             TidbClusterReady=True → IsTCReady(&tc) passes.
	defaultTC := makeTCWithRestartAnnotation(consts.ClusterPrefix, namespace, desiredVersion, true)

	// One TiDB pod on the new ASG node → areAllTiDBPodsOnNewNodeGroups returns true.
	tidbPod := getTestPod("db-tidb-0", namespace, string(cloudv1.ComponentTiDB), "node-tidb-preserve-no-extra")
	tidbNode := getTestNode("node-tidb-preserve-no-extra", newASG.Name, false) // self-managed: "Name" label

	eksScheme := runtime.NewScheme()
	g.Expect(corev1.AddToScheme(eksScheme)).Should(Succeed())
	g.Expect(opv1alpha1.AddToScheme(eksScheme)).Should(Succeed())

	eksCli := fake.NewClientBuilder().
		WithScheme(eksScheme).
		WithObjects(defaultTC, tidbPod, tidbNode).
		Build()

	u := &upgrader{
		infraCli:      infraCli,
		eventRecorder: record.NewFakeRecorder(100),
	}
	envProvider := &fakeEnvProvider{eksCli: eksCli}

	result, err := u.UpgradeTiDBNodes(ctx, uo, envProvider, cluster)

	// Restart annotations are pre-applied and zero-replica phases are skipped, so the upgrader
	// proceeds to UpdateClusterStatus in one reconcile.
	g.Expect(err).Should(BeNil(), "expected no error")
	g.Expect(result.Requeue).Should(BeFalse(),
		"expected function to proceed past readiness gate (no extra groups → new loop is no-op), "+
			"but got Requeue=true")
	g.Expect(result.RequeueAfter).Should(Equal(1*time.Second),
		"expected RequeueAfter=1s from UpdateClusterStatus, got %v", result.RequeueAfter)
}

// Test_UpgradeTiDBNodes_Preservation_AllExtraGroupsDeleted verifies that when the cluster has
// extra groups but all are deleted (Replicas=0), UpgradeTiDBNodes skips them all and proceeds
// past the readiness gate to UpdateClusterStatus, returning ctrl.Result{RequeueAfter: 1s}.
//
// This test PASSES on UNFIXED code — it establishes the preservation baseline for Bug 1.
//
// Setup: all restart mutations applied, all pods on new nodes, default TC ready,
//
//	extra groups present but all have TiDB.Replicas=0 (Deleted()==true).
//
// Expected: function proceeds past readiness gate → returns RequeueAfter=1s (NOT Requeue=true).
//
// **Validates: Requirements 3.2, 3.4**
func Test_UpgradeTiDBNodes_Preservation_AllExtraGroupsDeleted(t *testing.T) {
	const (
		desiredVersion = "1.29"
		clusterName    = "preserve-all-deleted"
	)
	namespace := cloudv1.ClusterPrefixId(clusterName)

	ctx := context.Background()
	g := NewWithT(t)

	// Cluster with one extra group that has Replicas=0 → Deleted()==true.
	cluster := &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name: clusterName,
			Labels: map[string]string{
				cloudv1.ProviderTypeLabel: consts.ProviderTypeEKS,
			},
		},
		Spec: cloudv1.ClusterSpec{
			Cluster: cloudv1.ClusterTopology{
				TiDB: cloudv1.TiDB{
					Replicas:     1,
					InstanceType: "m5.xlarge", // must match makeReadyASG's launch template instance type
				},
			},
			ExtraGroups: []cloudv1.ExtraGroupSpec{
				{
					Name: "deleted-group",
					TiDB: cloudv1.TiDB{
						Replicas: 0, // Deleted() == true
					},
				},
			},
		},
	}

	// UpgradeOperation with scaleOutExtraTiDB=false.
	uo := makeTestUpgradeOperation("preserve-uo-all-deleted", namespace, desiredVersion)

	// One new-version ASG.
	newASG := makeReadyASG("tidb-ng-preserve-all-deleted", clusterName, desiredVersion)

	infraScheme := runtime.NewScheme()
	g.Expect(v1beta1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(cloudv1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(ec2v1alpha1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())

	infraCli := fake.NewClientBuilder().
		WithScheme(infraScheme).
		WithObjects(newASG, makeReadyASGLaunchTemplate(newASG.Name), cluster).
		WithStatusSubresource(cluster).
		Build()

	// Default TC: restart annotation set, TidbClusterReady=True.
	defaultTC := makeTCWithRestartAnnotation(consts.ClusterPrefix, namespace, desiredVersion, true)

	// One TiDB pod on the new ASG node.
	tidbPod := getTestPod("db-tidb-0", namespace, string(cloudv1.ComponentTiDB), "node-tidb-preserve-deleted")
	tidbNode := getTestNode("node-tidb-preserve-deleted", newASG.Name, false)

	eksScheme := runtime.NewScheme()
	g.Expect(corev1.AddToScheme(eksScheme)).Should(Succeed())
	g.Expect(opv1alpha1.AddToScheme(eksScheme)).Should(Succeed())

	eksCli := fake.NewClientBuilder().
		WithScheme(eksScheme).
		WithObjects(defaultTC, tidbPod, tidbNode).
		Build()

	u := &upgrader{
		infraCli:      infraCli,
		eventRecorder: record.NewFakeRecorder(100),
	}
	envProvider := &fakeEnvProvider{eksCli: eksCli}

	result, err := u.UpgradeTiDBNodes(ctx, uo, envProvider, cluster)

	// All extra groups are deleted → both the restart loop and the readiness loop skip them.
	// The function should proceed to UpdateClusterStatus and return RequeueAfter=1s.
	g.Expect(err).Should(BeNil(), "expected no error")
	g.Expect(result.Requeue).Should(BeFalse(),
		"expected function to proceed past readiness gate (all extra groups deleted → skipped), "+
			"but got Requeue=true")
	g.Expect(result.RequeueAfter).Should(Equal(1*time.Second),
		"expected RequeueAfter=1s from UpdateClusterStatus, got %v", result.RequeueAfter)
}

// Test_UpgradeTiDBNodes_AllTCsReady_ProceedsToScaleIn verifies that when all TCs (default + extra
// groups) are ready, UpgradeTiDBNodes proceeds past the unified readiness gate to UpdateClusterStatus,
// returning ctrl.Result{RequeueAfter: 1s}.
//
// This is the "all TCs ready" baseline for Bug 1 — the fix must not block when all TCs are healthy.
//
// **Validates: Requirements 2.1, 2.2 (positive case)**
func Test_UpgradeTiDBNodes_AllTCsReady_ProceedsToScaleIn(t *testing.T) {
	const (
		desiredVersion = "1.29"
		clusterName    = "all-ready-cluster"
		extraGroupName = "analytics"
	)
	namespace := cloudv1.ClusterPrefixId(clusterName)

	ctx := context.Background()
	g := NewWithT(t)

	cluster := makeTestCluster(clusterName, desiredVersion, extraGroupName)

	extraGroupNamer := namer.NewNamer().
		ForTiDBGroup(cluster, &cluster.Spec.ExtraGroups[0], "").
		WithComponent(string(cloudv1.ComponentTiDB))
	extraGroupTCName := extraGroupNamer.GetTiDBName()

	// scaleOutExtraTiDB=false to skip the extra tidb scale-in path.
	uo := makeTestUpgradeOperation("all-ready-uo", namespace, desiredVersion)

	newASG := makeReadyASG("tidb-ng-all-ready", clusterName, desiredVersion)

	infraScheme := runtime.NewScheme()
	g.Expect(v1beta1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(cloudv1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(ec2v1alpha1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())

	infraCli := fake.NewClientBuilder().
		WithScheme(infraScheme).
		WithObjects(newASG, makeReadyASGLaunchTemplate(newASG.Name), cluster).
		WithStatusSubresource(cluster).
		Build()

	// Default TC: restart annotation set, TidbClusterReady=True.
	defaultTC := makeTCWithRestartAnnotation(consts.ClusterPrefix, namespace, desiredVersion, true)

	// Extra group TC: restart annotation set, TidbClusterReady=True (all TCs ready).
	extraGroupTC := makeTCWithRestartAnnotation(extraGroupTCName, namespace, desiredVersion, true)

	tidbPod := getTestPod("db-tidb-0", namespace, string(cloudv1.ComponentTiDB), "node-tidb-all-ready")
	tidbNode := getTestNode("node-tidb-all-ready", newASG.Name, false)

	eksScheme := runtime.NewScheme()
	g.Expect(corev1.AddToScheme(eksScheme)).Should(Succeed())
	g.Expect(opv1alpha1.AddToScheme(eksScheme)).Should(Succeed())

	eksCli := fake.NewClientBuilder().
		WithScheme(eksScheme).
		WithObjects(defaultTC, extraGroupTC, tidbPod, tidbNode).
		Build()

	u := &upgrader{
		infraCli:      infraCli,
		eventRecorder: record.NewFakeRecorder(100),
	}
	envProvider := &fakeEnvProvider{eksCli: eksCli}

	result, err := u.UpgradeTiDBNodes(ctx, uo, envProvider, cluster)

	// All TCs ready → proceeds to UpdateClusterStatus.
	g.Expect(err).Should(BeNil(), "expected no error")
	g.Expect(result.Requeue).Should(BeFalse(),
		"expected function to proceed past readiness gate (all TCs ready), but got Requeue=true")
	g.Expect(result.RequeueAfter).Should(Equal(1*time.Second),
		"expected RequeueAfter=1s from UpdateClusterStatus, got %v", result.RequeueAfter)
}

func Test_UpgradeTiDBNodes_ExtraGroup_RestartTiProxyFirst(t *testing.T) {
	const (
		desiredVersion = "1.29"
		clusterName    = "extra-group-tiproxy-first"
		extraGroupName = "analytics"
	)
	namespace := cloudv1.ClusterPrefixId(clusterName)

	ctx := context.Background()
	g := NewWithT(t)

	cluster := makeTestCluster(clusterName, desiredVersion, extraGroupName)
	uo := makeTestUpgradeOperation("uo-extra-group-tiproxy-first", namespace, desiredVersion)
	// Avoid infraCli status update dependency in this focused behavior test.
	// The upgrader will still return the same ctrl.Result and mutate TidbClusters,
	// but uo.UpdateReason becomes a no-op when reason is unchanged.
	uo.Status.Reason = ReasonRestartingTiProxy

	extraGroupNamer := namer.NewNamer().
		ForTiDBGroup(cluster, &cluster.Spec.ExtraGroups[0], "").
		WithComponent(string(cloudv1.ComponentTiDB))
	extraGroupTCName := extraGroupNamer.GetTiDBName()

	newASG := makeReadyASG("tidb-ng-extra-group-tiproxy-first", clusterName, desiredVersion)

	infraScheme := runtime.NewScheme()
	g.Expect(v1beta1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(cloudv1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(ec2v1alpha1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())

	infraCli := fake.NewClientBuilder().
		WithScheme(infraScheme).
		WithObjects(newASG, makeReadyASGLaunchTemplate(newASG.Name), cluster).
		WithStatusSubresource(cluster).
		Build()

	// Default + extra group both have TiProxy but no restart annotation yet.
	defaultTC := makeTCWithTiProxyNoRestart(consts.ClusterPrefix, namespace, true)
	extraGroupTC := makeTCWithTiProxyNoRestart(extraGroupTCName, namespace, true)

	tidbPod := getTestPod("db-tidb-0", namespace, string(cloudv1.ComponentTiDB), "node-extra-group-tiproxy-first")
	tidbNode := getTestNode("node-extra-group-tiproxy-first", newASG.Name, false)

	eksScheme := runtime.NewScheme()
	g.Expect(corev1.AddToScheme(eksScheme)).Should(Succeed())
	g.Expect(opv1alpha1.AddToScheme(eksScheme)).Should(Succeed())

	eksCli := fake.NewClientBuilder().
		WithScheme(eksScheme).
		WithObjects(defaultTC, extraGroupTC, tidbPod, tidbNode).
		Build()

	u := &upgrader{
		infraCli:      infraCli,
		eventRecorder: record.NewFakeRecorder(100),
	}
	envProvider := &fakeEnvProvider{eksCli: eksCli}

	result, err := u.UpgradeTiDBNodes(ctx, uo, envProvider, cluster)
	g.Expect(err).Should(BeNil(), "expected no error")
	g.Expect(result.RequeueAfter).Should(Equal(1*time.Minute), "expected to requeue after updating all TCs for TiProxy restart")

	updatedDefault := &opv1alpha1.TidbCluster{}
	g.Expect(eksCli.Get(ctx, client.ObjectKey{Name: consts.ClusterPrefix, Namespace: namespace}, updatedDefault)).Should(Succeed())

	g.Expect(updatedDefault.Spec.TiProxy).ShouldNot(BeNil())
	g.Expect(updatedDefault.Spec.TiProxy.ComponentSpec.Annotations).ShouldNot(BeNil())
	g.Expect(updatedDefault.Spec.TiProxy.ComponentSpec.Annotations[RestartTiDBAnnotation]).Should(Equal(desiredVersion))

	// Remaining components should NOT be restarted in the same reconcile.
	g.Expect(updatedDefault.Spec.PD.ComponentSpec.Annotations).Should(BeNil())
	g.Expect(updatedDefault.Spec.TiKV.ComponentSpec.Annotations).Should(BeNil())
	g.Expect(updatedDefault.Spec.TiDB.ComponentSpec.Annotations).Should(BeNil())

	updatedExtra := &opv1alpha1.TidbCluster{}
	g.Expect(eksCli.Get(ctx, client.ObjectKey{Name: extraGroupTCName, Namespace: namespace}, updatedExtra)).Should(Succeed())

	g.Expect(updatedExtra.Spec.TiProxy).ShouldNot(BeNil())
	g.Expect(updatedExtra.Spec.TiProxy.ComponentSpec.Annotations).ShouldNot(BeNil())
	g.Expect(updatedExtra.Spec.TiProxy.ComponentSpec.Annotations[RestartTiDBAnnotation]).Should(Equal(desiredVersion))

	// Remaining components should NOT be restarted in the same reconcile.
	g.Expect(updatedExtra.Spec.PD.ComponentSpec.Annotations).Should(BeNil())
	g.Expect(updatedExtra.Spec.TiKV.ComponentSpec.Annotations).Should(BeNil())
	g.Expect(updatedExtra.Spec.TiDB.ComponentSpec.Annotations).Should(BeNil())
}

// fakeExtraTiDB implements extratidb.Handler for unit tests.
// PrepareNodesV2 always returns nodegroupName; AreAllNodesReady always returns true.
type fakeExtraTiDB struct {
	nodegroupName string
}

func (f *fakeExtraTiDB) PrepareNodes(_ context.Context, _ client.Client, _ gitops.Provider, _ *cloudv1.Network, _ *cloudv1.Cluster) (string, error) {
	return f.nodegroupName, nil
}

func (f *fakeExtraTiDB) PrepareNodesV2(_ context.Context, _ client.Client, _ gitops.Provider, _ dataplane.EnvProvider, _ *cloudv1.Cluster) (string, error) {
	return f.nodegroupName, nil
}

func (f *fakeExtraTiDB) AreAllNodesReady(_ context.Context, _ client.Client, _ *cloudv1.Cluster) (bool, error) {
	return true, nil
}

func (f *fakeExtraTiDB) DeleteNodes(_ context.Context, _ client.Client, _ *cloudv1.Cluster) error {
	return nil
}

var _ extratidb.Handler = &fakeExtraTiDB{}

// Test_UpgradeTiDBNodes_ScaleOutExtraTiDB_ExtraGroupTCNotReady verifies Bug 1 when
// scaleOutExtraTiDB=true: after all restart mutations are applied, all pods are on new nodes
// (including one pod on the extra tidb nodegroup), and the default TC is ready, the function
// must wait for the extra group TC before proceeding to scale-in.
//
// Setup (scaleOutExtraTiDB=true, Bug Condition C1):
//   - scaleOutExtraTiDB=true → PrepareNodesV2 is called; extra tidb nodegroup = "extra-tidb-ng"
//   - All restart annotations already set → MutateTCForRestartingWithoutTiProxy and MutateTCForRestarting return false
//   - TiDB pods: one on new ASG node, one on extra tidb nodegroup node → areAllTiDBPodsOnNewNodeGroups returns true
//   - Default TC: TidbClusterReady=True
//   - Extra group TC: TidbClusterReady=False → should trigger requeue
//
// Expected (correct) behavior: return ctrl.Result{Requeue: true}
// Actual (unfixed) behavior: proceeds to scale-in path, returns error or RequeueAfter
//
// **Validates: Requirements 2.1, 2.2 (scaleOutExtraTiDB=true variant)**
func Test_UpgradeTiDBNodes_ScaleOutExtraTiDB_ExtraGroupTCNotReady(t *testing.T) {
	const (
		desiredVersion  = "1.29"
		clusterName     = "scaleout-bug-cluster"
		extraGroupName  = "analytics"
		extraTiDBNGName = "extra-tidb-ng"
	)
	namespace := cloudv1.ClusterPrefixId(clusterName)

	ctx := context.Background()
	g := NewWithT(t)

	cluster := makeTestCluster(clusterName, desiredVersion, extraGroupName)

	extraGroupNamer := namer.NewNamer().
		ForTiDBGroup(cluster, &cluster.Spec.ExtraGroups[0], "").
		WithComponent(string(cloudv1.ComponentTiDB))
	extraGroupTCName := extraGroupNamer.GetTiDBName()

	// scaleOutExtraTiDB=true (default when option is not "false")
	uo := &cloudv1.UpgradeOperation{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "scaleout-bug-uo",
			Namespace: namespace,
		},
		Status: cloudv1.UpgradeOperationStatus{
			CurrentStep: &cloudv1.Step{
				Version: desiredVersion,
				Options: map[string]string{
					ScaleOutExtraTiDBAnnotation: "true",
				},
			},
		},
	}

	newASG := makeReadyASG("tidb-ng-scaleout-bug", clusterName, desiredVersion)

	infraScheme := runtime.NewScheme()
	g.Expect(v1beta1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(cloudv1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(ec2v1alpha1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())

	infraCli := fake.NewClientBuilder().
		WithScheme(infraScheme).
		WithObjects(newASG, makeReadyASGLaunchTemplate(newASG.Name), cluster).
		WithStatusSubresource(cluster).
		Build()

	// Default TC: restart annotation set, TidbClusterReady=True.
	defaultTC := makeTCWithRestartAnnotation(consts.ClusterPrefix, namespace, desiredVersion, true)

	// Extra group TC: restart annotation set, TidbClusterReady=False → bug condition.
	extraGroupTC := makeTCWithRestartAnnotation(extraGroupTCName, namespace, desiredVersion, false)

	// Two TiDB pods:
	//   - one on the new ASG node (normal placement)
	//   - one on the extra tidb nodegroup node (scaleOutExtraTiDB path)
	tidbPod0 := getTestPod("db-tidb-0", namespace, string(cloudv1.ComponentTiDB), "node-tidb-scaleout-0")
	tidbPod1 := getTestPod("db-tidb-1", namespace, string(cloudv1.ComponentTiDB), "node-tidb-scaleout-extra")
	tidbNode0 := getTestNode("node-tidb-scaleout-0", newASG.Name, false)
	// Extra tidb node: its nodegroup label matches extraTiDBNGName so areAllTiDBPodsOnNewNodeGroups accepts it.
	tidbNodeExtra := getTestNode("node-tidb-scaleout-extra", extraTiDBNGName, false)

	eksScheme := runtime.NewScheme()
	g.Expect(corev1.AddToScheme(eksScheme)).Should(Succeed())
	g.Expect(opv1alpha1.AddToScheme(eksScheme)).Should(Succeed())

	eksCli := fake.NewClientBuilder().
		WithScheme(eksScheme).
		WithObjects(defaultTC, extraGroupTC, tidbPod0, tidbPod1, tidbNode0, tidbNodeExtra).
		Build()

	u := &upgrader{
		infraCli:      infraCli,
		extraTiDB:     &fakeExtraTiDB{nodegroupName: extraTiDBNGName},
		eventRecorder: record.NewFakeRecorder(100),
	}
	envProvider := &fakeEnvProvider{eksCli: eksCli}

	result, err := u.UpgradeTiDBNodes(ctx, uo, envProvider, cluster)

	// Fixed code: waits for extra group TC readiness before proceeding (Phase 7).
	g.Expect(err).Should(BeNil(), "expected no error from UpgradeTiDBNodes")
	g.Expect(result.Requeue).Should(BeFalse(),
		"UpgradeTiDBNodes (scaleOutExtraTiDB=true) should not proceed while extra group TC %q is not ready; got result=%+v",
		extraGroupTCName, result)
	g.Expect(result.RequeueAfter).Should(Equal(2*time.Minute),
		"expected RequeueAfter=2m while waiting for extra group TC readiness, got %v",
		result.RequeueAfter)
}

// Test_UpgradeTiDBNodes_ScaleOutExtraTiDB_AllTCsReady_ProceedsToScaleIn verifies that when
// scaleOutExtraTiDB=true and all TCs (default + extra groups) are ready, UpgradeTiDBNodes
// proceeds past the unified readiness gate into the scale-in path and returns
// ctrl.Result{RequeueAfter: 1s}.
//
// This is the positive/preservation case for the scaleOutExtraTiDB=true path.
//
// **Validates: Requirements 2.1, 2.2, 3.1 (scaleOutExtraTiDB=true positive case)**
func Test_UpgradeTiDBNodes_ScaleOutExtraTiDB_AllTCsReady_ProceedsToScaleIn(t *testing.T) {
	const (
		desiredVersion  = "1.29"
		clusterName     = "scaleout-allready-cluster"
		extraGroupName  = "analytics"
		extraTiDBNGName = "extra-tidb-ng"
	)
	namespace := cloudv1.ClusterPrefixId(clusterName)

	ctx := context.Background()
	g := NewWithT(t)

	cluster := makeTestCluster(clusterName, desiredVersion, extraGroupName)

	extraGroupNamer := namer.NewNamer().
		ForTiDBGroup(cluster, &cluster.Spec.ExtraGroups[0], "").
		WithComponent(string(cloudv1.ComponentTiDB))
	extraGroupTCName := extraGroupNamer.GetTiDBName()

	// scaleOutExtraTiDB=true
	uo := &cloudv1.UpgradeOperation{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "scaleout-allready-uo",
			Namespace: namespace,
		},
		Status: cloudv1.UpgradeOperationStatus{
			CurrentStep: &cloudv1.Step{
				Version: desiredVersion,
				Options: map[string]string{
					ScaleOutExtraTiDBAnnotation: "true",
				},
			},
		},
	}

	newASG := makeReadyASG("tidb-ng-scaleout-allready", clusterName, desiredVersion)

	infraScheme := runtime.NewScheme()
	g.Expect(v1beta1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(cloudv1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(ec2v1alpha1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())

	infraCli := fake.NewClientBuilder().
		WithScheme(infraScheme).
		WithObjects(newASG, makeReadyASGLaunchTemplate(newASG.Name), cluster).
		WithStatusSubresource(cluster).
		Build()

	// Default TC: restart annotation set, TidbClusterReady=True.
	// scaleOutExtraTiDB=true: the TC has been scaled out to 2 TiDB replicas (1 original + 1 extra).
	// Status fields are set so that IsComponentRestartedAndReady and the new ReadyReplicas gate pass.
	defaultTC := makeTCWithRestartAnnotation(consts.ClusterPrefix, namespace, desiredVersion, true)
	defaultTC.Spec.TiDB.Replicas = 2
	defaultTC.Status.TiDB.Phase = opv1alpha1.NormalPhase
	defaultTC.Status.TiDB.Members = map[string]opv1alpha1.TiDBMember{
		"db-tidb-0": {Name: "db-tidb-0", Health: true},
		"db-tidb-1": {Name: "db-tidb-1", Health: true},
	}
	defaultTC.Status.TiDB.StatefulSet = &apps.StatefulSetStatus{ReadyReplicas: 2}

	// Extra group TC: restart annotation set, TidbClusterReady=True (all TCs ready).
	extraGroupTC := makeTCWithRestartAnnotation(extraGroupTCName, namespace, desiredVersion, true)

	// Two TiDB pods: one on new ASG node, one on extra tidb nodegroup node.
	// The extra tidb pod name must contain a number for MutateTCToScaleInExtraTiDBPod to parse.
	tidbPod0 := getTestPod("db-tidb-0", namespace, string(cloudv1.ComponentTiDB), "node-tidb-allready-0")
	tidbPod1 := getTestPod("db-tidb-1", namespace, string(cloudv1.ComponentTiDB), "node-tidb-allready-extra")
	tidbNode0 := getTestNode("node-tidb-allready-0", newASG.Name, false)
	tidbNodeExtra := getTestNode("node-tidb-allready-extra", extraTiDBNGName, false)

	eksScheme := runtime.NewScheme()
	g.Expect(corev1.AddToScheme(eksScheme)).Should(Succeed())
	g.Expect(opv1alpha1.AddToScheme(eksScheme)).Should(Succeed())

	eksCli := fake.NewClientBuilder().
		WithScheme(eksScheme).
		WithObjects(defaultTC, extraGroupTC, tidbPod0, tidbPod1, tidbNode0, tidbNodeExtra).
		Build()

	u := &upgrader{
		infraCli:      infraCli,
		extraTiDB:     &fakeExtraTiDB{nodegroupName: extraTiDBNGName},
		eventRecorder: record.NewFakeRecorder(100),
	}
	envProvider := &fakeEnvProvider{eksCli: eksCli}

	result, err := u.UpgradeTiDBNodes(ctx, uo, envProvider, cluster)

	// All TCs ready → proceeds through scale-in path to UpdateClusterStatus.
	g.Expect(err).Should(BeNil(), "expected no error from UpgradeTiDBNodes")
	g.Expect(result.Requeue).Should(BeFalse(),
		"expected function to proceed past readiness gate (all TCs ready, scaleOutExtraTiDB=true), "+
			"but got Requeue=true; result=%+v", result)
	g.Expect(result.RequeueAfter).Should(Equal(1*time.Second),
		"expected RequeueAfter=1s from UpdateClusterStatus, got %v", result.RequeueAfter)
}

// Test_UpgradeTiDBNodes_StoragePhase_Blocked verifies that when the TiProxy phase is complete
// but the Storage phase annotations (PD/TiKV/TiFlash) have not been written yet, the upgrader
// writes the storage annotations and returns RequeueAfter=1m — it does NOT skip to TiDB.
//
// Setup:
//   - TiProxy: Replicas=0 (phase completes instantly — no mutation needed, no pod check, ready check skipped).
//   - PD: Replicas=1, annotation NOT set → triggers mutation.
//   - TiKV/TiFlash/TiDB: Replicas=0.
//
// Expected: Storage annotations mutated on TC, RequeueAfter=1m.
func Test_UpgradeTiDBNodes_StoragePhase_Blocked(t *testing.T) {
	const (
		desiredVersion = "1.29"
		clusterName    = "storage-blocked-cluster"
	)
	namespace := cloudv1.ClusterPrefixId(clusterName)
	ctx := context.Background()
	g := NewWithT(t)

	cluster := makeTestCluster(clusterName, desiredVersion, "")
	uo := makeTestUpgradeOperation("storage-blocked-uo", namespace, desiredVersion)
	// Pre-set Reason so uo.UpdateReason is a no-op (avoids fake-client update of uo status).
	uo.Status.Reason = ReasonRestartingStorage

	newASG := makeReadyASG("tidb-ng-storage-blocked", clusterName, desiredVersion)

	infraScheme := runtime.NewScheme()
	g.Expect(v1beta1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(cloudv1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(ec2v1alpha1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())

	infraCli := fake.NewClientBuilder().
		WithScheme(infraScheme).
		WithObjects(newASG, makeReadyASGLaunchTemplate(newASG.Name), cluster).
		WithStatusSubresource(cluster).
		Build()

	// TC: TiProxy phase already complete; PD Replicas=1 with NO annotation triggers Storage mutation.
	defaultTC := makeTCWithTiProxyPhaseComplete(consts.ClusterPrefix, namespace, desiredVersion)
	defaultTC.Spec.PD = &opv1alpha1.PDSpec{Replicas: 1}

	eksScheme := runtime.NewScheme()
	g.Expect(corev1.AddToScheme(eksScheme)).Should(Succeed())
	g.Expect(opv1alpha1.AddToScheme(eksScheme)).Should(Succeed())

	tiproxyObjs, _ := podsOnASGNode(namespace, newASG.Name, "-tiproxy", map[string]string{
		"db-tiproxy-0": string(cloudv1.ComponentTiProxy),
	})
	eksObjs := append([]client.Object{defaultTC}, tiproxyObjs...)

	eksCli := fake.NewClientBuilder().
		WithScheme(eksScheme).
		WithObjects(eksObjs...).
		Build()

	u := &upgrader{
		infraCli:      infraCli,
		eventRecorder: record.NewFakeRecorder(100),
	}
	envProvider := &fakeEnvProvider{eksCli: eksCli}

	result, err := u.UpgradeTiDBNodes(ctx, uo, envProvider, cluster)

	g.Expect(err).Should(BeNil())
	g.Expect(result.RequeueAfter).Should(Equal(1*time.Minute),
		"expected requeue after 1m while writing storage annotations, got %v", result.RequeueAfter)

	updatedTC := &opv1alpha1.TidbCluster{}
	g.Expect(eksCli.Get(ctx, client.ObjectKey{Name: consts.ClusterPrefix, Namespace: namespace}, updatedTC)).Should(Succeed())
	g.Expect(updatedTC.Spec.PD.ComponentSpec.Annotations[RestartTiDBAnnotation]).Should(Equal(desiredVersion),
		"PD restart annotation should have been written in the Storage phase")
	g.Expect(updatedTC.Spec.TiDB.ComponentSpec.Annotations).Should(BeNil(),
		"TiDB restart annotation must NOT be written before the Storage phase completes")
}

// Test_UpgradeTiDBNodes_TiDBPhase_RestartLast verifies that TiDB is restarted only after the
// TiProxy and Storage phases are complete. When TiProxy (Replicas=0) and storage components
// (Replicas=0) are all implicitly done, the upgrader targets TiDB next.
//
// Setup:
//   - TiProxy: Replicas=0 → Phase A completes instantly.
//   - PD/TiKV/TiFlash: Replicas=0 → Phase B completes instantly.
//   - TiDB: Replicas=1, annotation NOT set → triggers Phase C mutation.
//
// Expected: TiDB annotation written, RequeueAfter=1m.
func Test_UpgradeTiDBNodes_TiDBPhase_RestartLast(t *testing.T) {
	const (
		desiredVersion = "1.29"
		clusterName    = "tidb-last-cluster"
	)
	namespace := cloudv1.ClusterPrefixId(clusterName)
	ctx := context.Background()
	g := NewWithT(t)

	cluster := makeTestCluster(clusterName, desiredVersion, "")
	uo := makeTestUpgradeOperation("tidb-last-uo", namespace, desiredVersion)
	// Pre-set Reason so uo.UpdateReason is a no-op (avoids fake-client update of uo status).
	uo.Status.Reason = ReasonRestartingTiDB

	newASG := makeReadyASG("tidb-ng-tidb-last", clusterName, desiredVersion)

	infraScheme := runtime.NewScheme()
	g.Expect(v1beta1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(cloudv1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())
	g.Expect(ec2v1alpha1.SchemeBuilder.AddToScheme(infraScheme)).Should(Succeed())

	infraCli := fake.NewClientBuilder().
		WithScheme(infraScheme).
		WithObjects(newASG, makeReadyASGLaunchTemplate(newASG.Name), cluster, uo).
		WithStatusSubresource(cluster, uo).
		Build()
	defaultTC := makeTCWithTiProxyAllRestartAnnotation(consts.ClusterPrefix, namespace, desiredVersion, true)
	defaultTC.Spec.TiDB.ComponentSpec.Annotations = nil
	defaultTC.Spec.PD.Replicas = 1
	defaultTC.Status.PD = opv1alpha1.PDStatus{
		Phase:  opv1alpha1.NormalPhase,
		Synced: true,
		Members: map[string]opv1alpha1.PDMember{
			"pd-0": {Name: "pd-0", Health: true},
		},
	}

	eksScheme := runtime.NewScheme()
	g.Expect(corev1.AddToScheme(eksScheme)).Should(Succeed())
	g.Expect(opv1alpha1.AddToScheme(eksScheme)).Should(Succeed())

	componentObjs, _ := podsOnASGNode(namespace, newASG.Name, "-components", map[string]string{
		"db-tiproxy-0": string(cloudv1.ComponentTiProxy),
		"db-pd-0":      string(cloudv1.ComponentPD),
	})
	eksObjs := append([]client.Object{defaultTC}, componentObjs...)

	eksCli := fake.NewClientBuilder().
		WithScheme(eksScheme).
		WithObjects(eksObjs...).
		Build()

	u := &upgrader{
		infraCli:      infraCli,
		eventRecorder: record.NewFakeRecorder(100),
	}
	envProvider := &fakeEnvProvider{eksCli: eksCli}

	result, err := u.UpgradeTiDBNodes(ctx, uo, envProvider, cluster)

	g.Expect(err).Should(BeNil())
	g.Expect(result.RequeueAfter).Should(Equal(1*time.Minute),
		"expected requeue after 1m while writing TiDB annotation, got %v", result.RequeueAfter)

	updatedTC := &opv1alpha1.TidbCluster{}
	g.Expect(eksCli.Get(ctx, client.ObjectKey{Name: consts.ClusterPrefix, Namespace: namespace}, updatedTC)).Should(Succeed())
	g.Expect(updatedTC.Spec.TiDB.ComponentSpec.Annotations[RestartTiDBAnnotation]).Should(Equal(desiredVersion),
		"TiDB restart annotation should have been written in Phase C")
}