package generator

import (
	"context"
	"reflect"
	"testing"
	"time"

	"github.com/go-logr/zapr"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"go.uber.org/zap"

	entity "github.com/tidbcloud/devops/entity/k8supgrade"
	"github.com/tidbcloud/devops/pkg/infra/config"
)

const (
	testClusterID = "123456789"
)

func eksAdminNode(version string) *entity.Node {
	return &entity.Node{
		Kind:           string(cloudv1.AdminNodes),
		Version:        entity.NewMustVersion(version),
		MaxSkewVersion: entity.ManagedNodeGroupMaxSkewVersion,
	}
}

func makeTiDBClusterInfraInfo(id string, status cloudv1.ClusterPhase, tidbNodes, cdcNodes, dmNodes *entity.Node, extraNodes ...*entity.Node) *entity.TiDBClusterInfraInfo {
	ret := &entity.TiDBClusterInfraInfo{
		ID:             id,
		Status:         string(status),
		ComponentNodes: make([]*entity.Node, 0, 3),
		Topology:       map[string]int{"pd": 3, "tikv": 3, "tidb": 2},
	}
	if tidbNodes != nil {
		ret.ComponentNodes = append(ret.ComponentNodes, &entity.Node{
			Kind:           string(cloudv1.TiDBNodes),
			Count:          tidbNodes.Count,
			Version:        tidbNodes.Version,
			MaxSkewVersion: tidbNodes.MaxSkewVersion,
		})
	}
	if cdcNodes != nil {
		ret.ComponentNodes = append(ret.ComponentNodes, &entity.Node{
			Kind:           string(cloudv1.CDCNodes),
			Count:          cdcNodes.Count,
			Version:        cdcNodes.Version,
			MaxSkewVersion: cdcNodes.MaxSkewVersion,
		})
	}
	if dmNodes != nil {
		ret.ComponentNodes = append(ret.ComponentNodes, &entity.Node{
			Kind:           string(cloudv1.DMNodes),
			Count:          dmNodes.Count,
			Version:        dmNodes.Version,
			MaxSkewVersion: dmNodes.MaxSkewVersion,
		})
	}
	for _, node := range extraNodes {
		if node == nil {
			continue
		}
		ret.ComponentNodes = append(ret.ComponentNodes, &entity.Node{
			Kind:           node.Kind,
			Count:          node.Count,
			Version:        node.Version,
			MaxSkewVersion: node.MaxSkewVersion,
		})
	}
	return ret
}

func o11yNode(kind, version string) *entity.Node {
	return &entity.Node{
		Kind:           kind,
		Count:          1,
		Version:        entity.NewMustVersion(version),
		MaxSkewVersion: entity.ManagedNodeGroupMaxSkewVersion,
	}
}

func Test_generateUpgradeSteps(t *testing.T) {
	tests := []struct {
		name           string
		info           *entity.K8sInfo
		cfg            *config.RootConfig
		want           []*entity.UpgradeStep
		desiredVersion string
	}{
		// eks
		{
			name: "eks: no tidb clusters",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				ControlPlaneVersion: entity.NewMustVersion("1.29"),
				TiDBClusters:        []*entity.TiDBClusterInfraInfo{},
				AdminNodes:          []*entity.Node{eksAdminNode("1.29")},
			},
			cfg: &config.RootConfig{
				Services: config.Services{
					K8sUpgrade: config.K8sUpgradeServiceConfig{
						GeneratorConfig: config.K8sUpgradeGeneratorConfig{
							EKSAddonsConfig: map[string][]config.EKSAddon{
								"1.31": {{Name: "kube-proxy", Version: "v1.31.0-eksbuild.1"}},
							},
						},
					},
				},
			},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.30"},
				{Target: string(cloudv1.ControlPlane), Version: "1.31"},
				{Target: string(cloudv1.AdminNodes), Version: "1.31"},
				{Target: string(cloudv1.AddOns), Version: "1.31", Options: map[string]string{
					"kube-proxy": "v1.31.0-eksbuild.1",
				}},
			},
			desiredVersion: "1.31",
		},
		{
			name: "eks: a creating tidb cluster",
			info: &entity.K8sInfo{
				CloudProvider:        cloudv1.CloudProviderAWS,
				Is3LayerArch:         true,
				ControlPlaneVersion:  entity.NewMustVersion("1.29"),
				NumberOfTiDBClusters: 1,
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseCreating, nil, nil, nil),
				},
				AdminNodes: []*entity.Node{eksAdminNode("1.29")},
			},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.30"},
				{Target: string(cloudv1.ControlPlane), Version: "1.31"},
				{Target: string(cloudv1.AdminNodes), Version: "1.31"},
				{Target: "addons", Version: "1.31"},
			},
			desiredVersion: "1.31",
		},
		{
			name: "eks: all nodes already at control plane version",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				Is3LayerArch:        true,
				ControlPlaneVersion: entity.NewMustVersion("1.29"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.29"),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, nil, nil),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{eksAdminNode("1.29")},
			},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.30"},
				{Target: string(cloudv1.ControlPlane), Version: "1.31"},
				{Target: string(cloudv1.AdminNodes), Version: "1.31"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.31", ClusterId: testClusterID},
				{Target: "addons", Version: "1.31"},
			},
			desiredVersion: "1.31",
		},
		{
			name: "eks: control plane's version is newer than nodes' initially",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				Is3LayerArch:        true,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.26"),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, nil, nil),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{eksAdminNode("1.26")},
			},
			cfg: &config.RootConfig{},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.28"},
				{Target: string(cloudv1.AdminNodes), Version: "1.28"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.28", ClusterId: testClusterID},
				{Target: string(cloudv1.ControlPlane), Version: "1.29"},
				{Target: string(cloudv1.AdminNodes), Version: "1.29"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: "addons", Version: "1.29"},
			},
			desiredVersion: "1.29",
		},
		{
			name: "eks: use self-managed node groups",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				Is3LayerArch:        true,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, nil, nil),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{eksAdminNode("1.27")},
			},
			cfg: &config.RootConfig{},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.28"},
				{Target: string(cloudv1.ControlPlane), Version: "1.29"},
				{Target: string(cloudv1.AdminNodes), Version: "1.29"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: "addons", Version: "1.29"},
			},
			desiredVersion: "1.29",
		},
		{
			name: "eks: with cdc and dm nodes",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				Is3LayerArch:        true,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.ManagedNodeGroupMaxSkewVersion,
					}, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.ManagedNodeGroupMaxSkewVersion,
					}),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{eksAdminNode("1.27")},
			},
			cfg: &config.RootConfig{},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.28"},
				{Target: string(cloudv1.ControlPlane), Version: "1.29"},
				{Target: string(cloudv1.AdminNodes), Version: "1.29"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: string(cloudv1.CDCNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: string(cloudv1.DMNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: "addons", Version: "1.29"},
				{Target: string(cloudv1.EmptyNodePools), Version: "1.29", Options: map[string]string{"delay": "10m"}}},
			desiredVersion: "1.29",
		},
		{
			name: "eks: with o11y k8s and tidb nodes",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				Is3LayerArch:        true,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, nil, nil,
						o11yNode(o11yTiDBNodes, "1.27"),
					),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{eksAdminNode("1.27")},
				O11YK8SNodes:         []*entity.Node{o11yNode(o11yK8sNodes, "1.27")},
			},
			cfg: &config.RootConfig{},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.28"},
				{Target: string(cloudv1.ControlPlane), Version: "1.29"},
				{Target: string(cloudv1.AdminNodes), Version: "1.29"},
				{Target: o11yK8sNodes, Version: "1.29"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: o11yTiDBNodes, Version: "1.29", ClusterId: testClusterID},
				{Target: "addons", Version: "1.29"},
			},
			desiredVersion: "1.29",
		},
		{
			name: "eks: generates only one o11y k8s step for multiple tidb clusters",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				Is3LayerArch:        true,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, nil, nil),
					makeTiDBClusterInfraInfo("another-cluster", cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, nil, nil),
				},
				NumberOfTiDBClusters: 2,
				AdminNodes:           []*entity.Node{eksAdminNode("1.27")},
				O11YK8SNodes:         []*entity.Node{o11yNode(o11yK8sNodes, "1.27")},
			},
			cfg: &config.RootConfig{},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.28"},
				{Target: string(cloudv1.ControlPlane), Version: "1.29"},
				{Target: string(cloudv1.AdminNodes), Version: "1.29"},
				{Target: o11yK8sNodes, Version: "1.29"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: "another-cluster"},
				{Target: "addons", Version: "1.29"},
			},
			desiredVersion: "1.29",
		},
		{
			name: "eks: skips o11y nodes already at generated target version",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				Is3LayerArch:        true,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, nil, nil,
						o11yNode(o11yTiDBNodes, "1.29"),
					),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{eksAdminNode("1.27")},
				O11YK8SNodes:         []*entity.Node{o11yNode(o11yK8sNodes, "1.29")},
			},
			cfg: &config.RootConfig{},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.28"},
				{Target: string(cloudv1.ControlPlane), Version: "1.29"},
				{Target: string(cloudv1.AdminNodes), Version: "1.29"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: "addons", Version: "1.29"},
			},
			desiredVersion: "1.29",
		},
		{
			name: "eks: max skew version 0 should fallback to 2",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				Is3LayerArch:        true,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: 0,
					}, nil, nil),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes: []*entity.Node{{
					Kind:           string(cloudv1.AdminNodes),
					Version:        entity.NewMustVersion("1.27"),
					MaxSkewVersion: 0,
				}},
			},
			cfg: &config.RootConfig{},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.28"},
				{Target: string(cloudv1.ControlPlane), Version: "1.29"},
				{Target: string(cloudv1.AdminNodes), Version: "1.29"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: "addons", Version: "1.29"},
			},
			desiredVersion: "1.29",
		},
		{
			name: "eks: skip empty kind and do not generate empty target",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				Is3LayerArch:        true,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					{
						ID:     testClusterID,
						Status: string(cloudv1.ClusterPhaseRunning),
						ComponentNodes: []*entity.Node{
							{
								Kind:           "",
								Version:        entity.NewMustVersion("1.27"),
								MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
							},
							{
								Kind:           string(cloudv1.TiDBNodes),
								Version:        entity.NewMustVersion("1.27"),
								MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
							},
						},
						Topology: map[string]int{"pd": 3, "tikv": 3, "tidb": 2},
					},
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{eksAdminNode("1.27")},
			},
			cfg: &config.RootConfig{},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.28"},
				{Target: string(cloudv1.ControlPlane), Version: "1.29"},
				{Target: string(cloudv1.AdminNodes), Version: "1.29"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: "addons", Version: "1.29"},
			},
			desiredVersion: "1.29",
		},
		{
			name: "eks: two tidb clusters in a k8s",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				Is3LayerArch:        true,
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, nil, nil),
					makeTiDBClusterInfraInfo("another-cluster", cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, nil, nil),
				},
				NumberOfTiDBClusters: 2,
				AdminNodes:           []*entity.Node{eksAdminNode("1.27")},
			},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.28"},
				{Target: string(cloudv1.ControlPlane), Version: "1.29"},
				{Target: string(cloudv1.AdminNodes), Version: "1.29"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: "another-cluster"},
				{Target: "addons", Version: "1.29"}},
			desiredVersion: "1.29",
		},
		{
			name: "eks: upgrade addons",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				ControlPlaneVersion: entity.NewMustVersion("1.28"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.28"),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, nil, nil),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{eksAdminNode("1.28")},
			},
			cfg: &config.RootConfig{
				Services: config.Services{
					K8sUpgrade: config.K8sUpgradeServiceConfig{
						GeneratorConfig: config.K8sUpgradeGeneratorConfig{
							EKSAddonsConfig: map[string][]config.EKSAddon{
								"1.30": {{Name: "kube-proxy", Version: "cloudv1.30.0-eksbuild.3"}},
							},
						},
					},
				},
			},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.29"},
				{Target: string(cloudv1.ControlPlane), Version: "1.30"},
				{Target: string(cloudv1.AdminNodes), Version: "1.30"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.30", ClusterId: testClusterID},
				{Target: string(cloudv1.AddOns), Version: "1.30", Options: map[string]string{"kube-proxy": "cloudv1.30.0-eksbuild.3"}},
			},
			desiredVersion: "1.30",
		},
		// gke
		{
			name: "gke: common case",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderGCP,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.NodePoolMaxSkewVersion,
					}, nil, nil),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{{Kind: string(cloudv1.AdminNodes), Version: entity.NewMustVersion("1.27"), MaxSkewVersion: entity.NodePoolMaxSkewVersion}},
			},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.28"},
				{Target: string(cloudv1.ControlPlane), Version: "1.29"},
				{Target: string(cloudv1.AdminNodes), Version: "1.29"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: string(cloudv1.EmptyNodePools), Version: "1.29", Options: map[string]string{"delay": "10m"}},
			},
			desiredVersion: "1.29",
		},
		{
			name: "gke: control plane's version is newer than nodes' initially",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderGCP,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.NodePoolMaxSkewVersion,
					}, nil, nil),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{{Kind: string(cloudv1.AdminNodes), Version: entity.NewMustVersion("1.26"), MaxSkewVersion: entity.NodePoolMaxSkewVersion}},
			},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.28"},
				{Target: string(cloudv1.AdminNodes), Version: "1.28"},
				{Target: string(cloudv1.ControlPlane), Version: "1.29"},
				{Target: string(cloudv1.AdminNodes), Version: "1.29"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: string(cloudv1.EmptyNodePools), Version: "1.29", Options: map[string]string{"delay": "10m"}},
			},
			desiredVersion: "1.29",
		},
		{
			name: "gke: two tidb clusters",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderGCP,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.NodePoolMaxSkewVersion,
					}, nil, nil),
					makeTiDBClusterInfraInfo("second-cluster", cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.NodePoolMaxSkewVersion,
					}, nil, nil),
				},
				NumberOfTiDBClusters: 2,
				AdminNodes:           []*entity.Node{{Kind: string(cloudv1.AdminNodes), Version: entity.NewMustVersion("1.27"), MaxSkewVersion: entity.NodePoolMaxSkewVersion}},
			},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.28"},
				{Target: string(cloudv1.ControlPlane), Version: "1.29"},
				{Target: string(cloudv1.AdminNodes), Version: "1.29"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: "second-cluster"},
				{Target: string(cloudv1.EmptyNodePools), Version: "1.29", Options: map[string]string{"delay": "10m"}},
			},
			desiredVersion: "1.29",
		},
		{
			name: "gke: upgrades o11y nodes",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderGCP,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.NodePoolMaxSkewVersion,
					}, nil, nil,
						o11yNode(o11yTiDBNodes, "1.27"),
					),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{{Kind: string(cloudv1.AdminNodes), Version: entity.NewMustVersion("1.27"), MaxSkewVersion: entity.NodePoolMaxSkewVersion}},
				O11YK8SNodes:         []*entity.Node{o11yNode(o11yK8sNodes, "1.27")},
			},
			want: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane), Version: "1.28"},
				{Target: string(cloudv1.ControlPlane), Version: "1.29"},
				{Target: string(cloudv1.AdminNodes), Version: "1.29"},
				{Target: o11yK8sNodes, Version: "1.29"},
				{Target: string(cloudv1.TiDBNodes), Version: "1.29", ClusterId: testClusterID},
				{Target: o11yTiDBNodes, Version: "1.29", ClusterId: testClusterID},
				{Target: string(cloudv1.EmptyNodePools), Version: "1.29", Options: map[string]string{"delay": "10m"}},
			},
			desiredVersion: "1.29",
		},
		// test invalid versions
		{
			name: "all versions are invalid",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				ControlPlaneVersion: entity.NewMustVersion(""),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion(""),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, nil, nil),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{eksAdminNode("")},
			},
			want:           nil,
			desiredVersion: "",
		},
		{
			name: "some versions are invalid",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion(""),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, nil, nil),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{eksAdminNode("1.27")},
			},
			want:           nil,
			desiredVersion: "",
		},
		{
			name: "control plane version is invalid",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				ControlPlaneVersion: entity.NewMustVersion(""),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, nil, nil),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{eksAdminNode("1.27")},
			},
			want:           nil,
			desiredVersion: "",
		},
		{
			name: "admin nodes version is invalid",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion("1.27"),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, nil, nil),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{eksAdminNode("")},
			},
			want:           nil,
			desiredVersion: "",
		},
		{
			name: "tidb cluster nodes version is invalid",
			info: &entity.K8sInfo{
				CloudProvider:       cloudv1.CloudProviderAWS,
				ControlPlaneVersion: entity.NewMustVersion("1.27"),
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{
						Version:        entity.NewMustVersion(""),
						MaxSkewVersion: entity.SelfManagedNodeGroupMaxSkewVersion,
					}, &entity.Node{
						Version:        entity.NewMustVersion(""),
						MaxSkewVersion: entity.ManagedNodeGroupMaxSkewVersion,
					}, nil),
				},
				NumberOfTiDBClusters: 1,
				AdminNodes:           []*entity.Node{eksAdminNode("1.27")},
			},
			want:           nil,
			desiredVersion: "",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			zapLog, _ := zap.NewDevelopment()
			var g gen
			if tt.cfg != nil {
				g = gen{logger: zapr.NewLogger(zapLog), cfg: tt.cfg}
			} else {
				g = gen{logger: zapr.NewLogger(zapLog)}
			}
			gotSteps, gotDesiredVersion := g.generateUpgradeSteps(context.TODO(), tt.info)
			if !reflect.DeepEqual(gotSteps, tt.want) || gotDesiredVersion != tt.desiredVersion {
				t.Errorf("generateUpgradeSteps() = %v, %v want %v, %v", gotSteps, gotDesiredVersion, tt.want, tt.desiredVersion)
			}
		})
	}
}

func Test_estimateDuration(t *testing.T) {
	tests := []struct {
		name  string
		info  *entity.K8sInfo
		steps []*entity.UpgradeStep
		want  time.Duration
	}{
		{
			name: "upgrade eks without dm and cdc",
			info: &entity.K8sInfo{
				CloudProvider: cloudv1.CloudProviderAWS,
				AdminNodes:    []*entity.Node{{Kind: string(cloudv1.AdminNodes), Count: 3}},
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{}, nil, nil),
				},
			},
			steps: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane)},                        // 15 mins
				{Target: string(cloudv1.AdminNodes)},                          // 15 mins
				{Target: string(cloudv1.TiDBNodes), ClusterId: testClusterID}, // 45 mins
			},
			want: 75 * time.Minute,
		},
		{
			name: "upgrade eks with dm and cdc",
			info: &entity.K8sInfo{
				CloudProvider: cloudv1.CloudProviderAWS,
				AdminNodes:    []*entity.Node{{Kind: "admin-nodes", Count: 3}},
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{}, &entity.Node{Count: 3}, &entity.Node{Count: 2}),
				},
			},
			steps: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane)},                        // 15 mins
				{Target: string(cloudv1.AdminNodes)},                          // 15 mins
				{Target: string(cloudv1.TiDBNodes), ClusterId: testClusterID}, // 45 mins
				{Target: string(cloudv1.CDCNodes), ClusterId: testClusterID},  // 15 mins
				{Target: string(cloudv1.DMNodes), ClusterId: testClusterID},   // 10 mins
			},
			want: 100 * time.Minute,
		},
		{
			name: "upgrade eks with o11y nodes",
			info: &entity.K8sInfo{
				CloudProvider: cloudv1.CloudProviderAWS,
				AdminNodes:    []*entity.Node{{Kind: string(cloudv1.AdminNodes), Count: 3}},
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{}, nil, nil),
				},
			},
			steps: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane)},                        // 15 mins
				{Target: string(cloudv1.AdminNodes)},                          // 15 mins
				{Target: string(cloudv1.TiDBNodes), ClusterId: testClusterID}, // 45 mins
				{Target: o11yK8sNodes},                                        // 10 mins
				{Target: o11yTiDBNodes, ClusterId: testClusterID},             // 10 mins
			},
			want: 95 * time.Minute,
		},
		{
			name: "upgrade gke without dm and cdc",
			info: &entity.K8sInfo{
				CloudProvider: cloudv1.CloudProviderGCP,
				AdminNodes:    []*entity.Node{{Kind: string(cloudv1.AdminNodes), Count: 2}},
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{}, nil, nil),
				},
			},
			steps: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane)},                                               // 35 mins
				{Target: string(cloudv1.AdminNodes)},                                                 // 10 mins
				{Target: string(cloudv1.TiDBNodes), ClusterId: testClusterID},                        // 45 mins
				{Target: string(cloudv1.EmptyNodePools), Options: map[string]string{"delay": "10m"}}, // 20 mins
			},
			want: 110 * time.Minute,
		},
		{
			name: "empty node pools with custom delay",
			info: &entity.K8sInfo{
				CloudProvider: cloudv1.CloudProviderGCP,
			},
			steps: []*entity.UpgradeStep{
				{Target: string(cloudv1.EmptyNodePools), Options: map[string]string{"delay": "1h2m3s"}},
			},
			want: 1*time.Hour + 12*time.Minute + 3*time.Second,
		},
		{
			name: "upgrade gke with cdc and dm",
			info: &entity.K8sInfo{
				CloudProvider: cloudv1.CloudProviderGCP,
				AdminNodes:    []*entity.Node{{Kind: string(cloudv1.AdminNodes), Count: 4}},
				TiDBClusters: []*entity.TiDBClusterInfraInfo{
					makeTiDBClusterInfraInfo(testClusterID, cloudv1.ClusterPhaseRunning, &entity.Node{}, &entity.Node{Count: 1}, &entity.Node{Count: 2}),
				},
			},
			steps: []*entity.UpgradeStep{
				{Target: string(cloudv1.ControlPlane)},                        // 35 mins
				{Target: string(cloudv1.AdminNodes)},                          // 20 mins
				{Target: string(cloudv1.TiDBNodes), ClusterId: testClusterID}, // 45 mins
				{Target: string(cloudv1.CDCNodes), ClusterId: testClusterID},  // 5 mins
				{Target: string(cloudv1.DMNodes), ClusterId: testClusterID},   // 10 mins
			},
			want: 115 * time.Minute,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := estimateDuration(tt.info, tt.steps); got != tt.want {
				t.Errorf("estimateDuration() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestGetDesiredVersion(t *testing.T) {
	tests := []struct {
		name string
		uo   *entity.UpgradeOperation
		want string
	}{
		{
			name: "uo is nil",
			uo:   nil,
		},
		{
			name: "no steps",
			uo:   &entity.UpgradeOperation{},
		},
		{
			name: "several steps",
			uo: &entity.UpgradeOperation{
				Steps: []*entity.UpgradeStep{
					{
						Target:  string(cloudv1.ControlPlane),
						Version: "1.29",
					},
					{
						Target:  string(cloudv1.ControlPlane),
						Version: "1.30",
					},
				},
			},
			want: "1.30",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := GetDesiredVersion(tt.uo); got != tt.want {
				t.Errorf("GetDesiredVersion() = %v, want %v", got, tt.want)
			}
		})
	}
}

func Test_gen_generateStepForEKSAddons(t *testing.T) {
	type fields struct {
		cfg *config.RootConfig
	}
	type args struct {
		info            *entity.K8sInfo
		controlPlaneVer entity.Version
	}
	tests := []struct {
		name   string
		fields fields
		args   args
		want   *entity.UpgradeStep
	}{
		{
			name: "needUpgradeEKSAddons returns false because cloud provider is not AWS",
			fields: fields{
				cfg: &config.RootConfig{},
			},
			args: args{
				info: &entity.K8sInfo{
					CloudProvider: cloudv1.CloudProviderGCP,
				},
				controlPlaneVer: entity.NewMustVersion("1.29"),
			},
			want: nil,
		},
		{
			name: "3-layer architecture, should return addons step",
			fields: fields{
				cfg: &config.RootConfig{},
			},
			args: args{
				info: &entity.K8sInfo{
					CloudProvider: cloudv1.CloudProviderAWS,
					Is3LayerArch:  true,
				},
				controlPlaneVer: entity.NewMustVersion("1.30"),
			},
			want: &entity.UpgradeStep{
				Target:  "addons",
				Version: "1.30",
			},
		},
		{
			name: "Non-3-layer, addons config exists",
			fields: fields{
				cfg: &config.RootConfig{
					Services: config.Services{
						K8sUpgrade: config.K8sUpgradeServiceConfig{
							GeneratorConfig: config.K8sUpgradeGeneratorConfig{
								EKSAddonsConfig: map[string][]config.EKSAddon{
									"1.30": {
										{Name: "kube-proxy", Version: "v1.30.0-eksbuild.1"},
										{Name: "vpc-cni", Version: "v1.15.0-eksbuild.1"},
									},
								},
							},
						},
					},
				},
			},
			args: args{
				info: &entity.K8sInfo{
					CloudProvider:        cloudv1.CloudProviderAWS,
					Is3LayerArch:         false,
					NumberOfTiDBClusters: 1, // To make needUpgradeEKSAddons check TiDBNodes
				},
				controlPlaneVer: entity.NewMustVersion("1.30"),
			},
			want: &entity.UpgradeStep{
				Target:  string(cloudv1.AddOns),
				Version: "1.30",
				Options: map[string]string{
					"kube-proxy": "v1.30.0-eksbuild.1",
					"vpc-cni":    "v1.15.0-eksbuild.1",
				},
			},
		},
		{
			name: "Non-3-layer, addons config does not exist for controlPlaneVer",
			fields: fields{
				cfg: &config.RootConfig{
					Services: config.Services{
						K8sUpgrade: config.K8sUpgradeServiceConfig{
							GeneratorConfig: config.K8sUpgradeGeneratorConfig{
								EKSAddonsConfig: map[string][]config.EKSAddon{
									"1.29": {{Name: "kube-proxy", Version: "v1.29.0-eksbuild.1"}},
								},
							},
						},
					},
				},
			},
			args: args{
				info: &entity.K8sInfo{
					CloudProvider:        cloudv1.CloudProviderAWS,
					Is3LayerArch:         false,
					NumberOfTiDBClusters: 1,
				},
				controlPlaneVer: entity.NewMustVersion("1.30"),
			},
			want: nil,
		},
		{
			name: "Non-3-layer, addons config exists but is empty for controlPlaneVer",
			fields: fields{
				cfg: &config.RootConfig{
					Services: config.Services{
						K8sUpgrade: config.K8sUpgradeServiceConfig{
							GeneratorConfig: config.K8sUpgradeGeneratorConfig{
								EKSAddonsConfig: map[string][]config.EKSAddon{
									"1.30": {}, // Empty addons
								},
							},
						},
					},
				},
			},
			args: args{
				info: &entity.K8sInfo{
					CloudProvider:        cloudv1.CloudProviderAWS,
					Is3LayerArch:         false,
					NumberOfTiDBClusters: 1,
				},
				controlPlaneVer: entity.NewMustVersion("1.30"),
			},
			want: nil,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			zapLog, _ := zap.NewDevelopment()
			g := &gen{
				logger: zapr.NewLogger(zapLog),
				cfg:    tt.fields.cfg,
			}
			if got := g.generateStepForEKSAddons(tt.args.info, tt.args.controlPlaneVer); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("gen.generateStepForEKSAddons() = %v, want %v", got, tt.want)
			}
		})
	}
}