//go:build unit || all
// +build unit all

package task

import (
	"testing"

	"github.com/onsi/gomega"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	apiv1 "github.com/tidbcloud/eks-provider/api/v1alpha1"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/config"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/namer"
)

func TestNodeLabelBuilders(t *testing.T) {
	tests := map[string]struct {
		provider            string
		expectedMasterLabel map[string]string
		expectedWorkerLabel map[string]string
	}{
		"AWS node labels should match expected map": {
			provider: consts.ProviderTypeAWS,
			expectedMasterLabel: map[string]string{
				"cloud.pingcap.com/provider":         consts.ProviderTypeAWS,
				cloudv1.RegionLabel:                  "us-west-2",
				consts.TenantLabel:                   "tenant-1",
				cloudv1.ProjectLabel:                 "project-1",
				"cloud.pingcap.com/component":        ComponentName,
				consts.TagKeyComponent:               ComponentName,
				consts.TagKeyUsedBy:                  "customer-tidb",
				namer.DMJobNameLabelKey:              "dmjob-1",
				namer.ClusterSelectorKey:             "cluster-1",
				namer.ComponentSelectorKey:           ComponentName,
				namer.NamespaceSelectorKey:           "tc-cluster-1",
				"tidbcloud-ngset-owner":              "d677252c",
				consts.GardenerWorkerPrefixLabelName: "cluster-1-dm",
				namer.DMMemberSpecLabelKey:           "master-spec",
				namer.DMMasterLabelKey:               "true",
			},
			expectedWorkerLabel: map[string]string{
				"cloud.pingcap.com/provider":         consts.ProviderTypeAWS,
				cloudv1.RegionLabel:                  "us-west-2",
				consts.TenantLabel:                   "tenant-1",
				cloudv1.ProjectLabel:                 "project-1",
				"cloud.pingcap.com/component":        ComponentName,
				consts.TagKeyComponent:               ComponentName,
				consts.TagKeyUsedBy:                  "customer-tidb",
				namer.DMJobNameLabelKey:              "dmjob-1",
				namer.ClusterSelectorKey:             "cluster-1",
				namer.ComponentSelectorKey:           ComponentName,
				namer.NamespaceSelectorKey:           "tc-cluster-1",
				"tidbcloud-ngset-owner":              "d677252c",
				consts.GardenerWorkerPrefixLabelName: "cluster-1-dm",
				namer.DMMemberSpecLabelKey:           "worker-spec",
				namer.DMWorkerLabelKey:               "true",
			},
		},
		"EKS node labels should match expected map": {
			provider: consts.ProviderTypeEKS,
			expectedMasterLabel: map[string]string{
				"cloud.pingcap.com/provider":         consts.ProviderTypeEKS,
				cloudv1.RegionLabel:                  "us-west-2",
				consts.TenantLabel:                   "tenant-1",
				cloudv1.ProjectLabel:                 "project-1",
				"cloud.pingcap.com/component":        ComponentName,
				consts.TagKeyComponent:               ComponentName,
				consts.TagKeyUsedBy:                  "customer-tidb",
				namer.DMJobNameLabelKey:              "dmjob-1",
				namer.ClusterSelectorKey:             "cluster-1",
				namer.ComponentSelectorKey:           ComponentName,
				namer.NamespaceSelectorKey:           "tc-cluster-1",
				"tidbcloud-ngset-owner":              "d677252c",
				consts.GardenerWorkerPrefixLabelName: "cluster-1-dm",
				namer.DMMemberSpecLabelKey:           "master-spec",
				namer.DMMasterLabelKey:               "true",
			},
			expectedWorkerLabel: map[string]string{
				"cloud.pingcap.com/provider":         consts.ProviderTypeEKS,
				cloudv1.RegionLabel:                  "us-west-2",
				consts.TenantLabel:                   "tenant-1",
				cloudv1.ProjectLabel:                 "project-1",
				"cloud.pingcap.com/component":        ComponentName,
				consts.TagKeyComponent:               ComponentName,
				consts.TagKeyUsedBy:                  "customer-tidb",
				namer.DMJobNameLabelKey:              "dmjob-1",
				namer.ClusterSelectorKey:             "cluster-1",
				namer.ComponentSelectorKey:           ComponentName,
				namer.NamespaceSelectorKey:           "tc-cluster-1",
				"tidbcloud-ngset-owner":              "d677252c",
				consts.GardenerWorkerPrefixLabelName: "cluster-1-dm",
				namer.DMMemberSpecLabelKey:           "worker-spec",
				namer.DMWorkerLabelKey:               "true",
			},
		},
		"GKE node labels should match expected map": {
			provider: consts.ProviderTypeGKE,
			expectedMasterLabel: map[string]string{
				"cloud.pingcap.com/provider":         consts.ProviderTypeGKE,
				cloudv1.RegionLabel:                  "us-west-2",
				consts.TenantLabel:                   "tenant-1",
				cloudv1.ProjectLabel:                 "project-1",
				"cloud.pingcap.com/component":        ComponentName,
				consts.TagKeyComponent:               ComponentName,
				consts.TagKeyUsedBy:                  "customer-tidb",
				namer.DMJobNameLabelKey:              "dmjob-1",
				namer.ClusterSelectorKey:             "cluster-1",
				namer.ComponentSelectorKey:           ComponentName,
				namer.NamespaceSelectorKey:           "tc-cluster-1",
				"tidbcloud-ngset-owner":              "d677252c",
				consts.GardenerWorkerPrefixLabelName: "cluster-1-dm",
				namer.DMMemberSpecLabelKey:           "master-spec",
				namer.DMMasterLabelKey:               "true",
			},
			expectedWorkerLabel: map[string]string{
				"cloud.pingcap.com/provider":         consts.ProviderTypeGKE,
				cloudv1.RegionLabel:                  "us-west-2",
				consts.TenantLabel:                   "tenant-1",
				cloudv1.ProjectLabel:                 "project-1",
				"cloud.pingcap.com/component":        ComponentName,
				consts.TagKeyComponent:               ComponentName,
				consts.TagKeyUsedBy:                  "customer-tidb",
				namer.DMJobNameLabelKey:              "dmjob-1",
				namer.ClusterSelectorKey:             "cluster-1",
				namer.ComponentSelectorKey:           ComponentName,
				namer.NamespaceSelectorKey:           "tc-cluster-1",
				"tidbcloud-ngset-owner":              "d677252c",
				consts.GardenerWorkerPrefixLabelName: "cluster-1-dm",
				namer.DMMemberSpecLabelKey:           "worker-spec",
				namer.DMWorkerLabelKey:               "true",
			},
		},
		"AKS node labels should match expected map": {
			provider: consts.ProviderTypeAKS,
			expectedMasterLabel: map[string]string{
				"cloud.pingcap.com/provider-type":    consts.ProviderTypeAKS,
				cloudv1.RegionLabel:                  "us-west-2",
				cloudv1.ProjectLabel:                 "project-1",
				consts.TagKeyComponent:               ComponentName,
				consts.TagKeyUsedBy:                  "customer-tidb",
				namer.DMJobNameLabelKey:              "dmjob-1",
				namer.ClusterSelectorKey:             "cluster-1",
				namer.ComponentSelectorKey:           ComponentName,
				namer.NamespaceSelectorKey:           "tc-cluster-1",
				"tidbcloud-ngset-owner":              "d677252c",
				consts.GardenerWorkerPrefixLabelName: "cluster-1-dm",
				namer.DMMemberSpecLabelKey:           "master-spec",
				namer.DMMasterLabelKey:               "true",
			},
			expectedWorkerLabel: map[string]string{
				"cloud.pingcap.com/provider-type":    consts.ProviderTypeAKS,
				cloudv1.RegionLabel:                  "us-west-2",
				cloudv1.ProjectLabel:                 "project-1",
				consts.TagKeyComponent:               ComponentName,
				consts.TagKeyUsedBy:                  "customer-tidb",
				namer.DMJobNameLabelKey:              "dmjob-1",
				namer.ClusterSelectorKey:             "cluster-1",
				namer.ComponentSelectorKey:           ComponentName,
				namer.NamespaceSelectorKey:           "tc-cluster-1",
				"tidbcloud-ngset-owner":              "d677252c",
				consts.GardenerWorkerPrefixLabelName: "cluster-1-dm",
				namer.DMMemberSpecLabelKey:           "worker-spec",
				namer.DMWorkerLabelKey:               "true",
			},
		},
	}

	for name, tt := range tests {
		t.Run(name, func(t *testing.T) {
			g := gomega.NewWithT(t)

			baseTask := &BaseTask{
				Context: &Context{
					Namer: namer.NewNamer().WithClusterName("cluster-1").WithTiDBNamespace("tc-cluster-1"),
					Config: config.RootConfig{
						ControllerConfig: config.ControllerConfig{
							Provider: tt.provider,
							Network: config.NetworkConfig{
								Region: "us-west-2",
							},
						},
					},
					TenantID:  "tenant-1",
					ProjectID: "project-1",
					DMJob: &apiv1.DMJob{
						ObjectMeta: metav1.ObjectMeta{
							Name:   "dmjob-1",
							Labels: map[string]string{},
						},
						Spec: apiv1.DMJobSpec{
							ClusterName: "cluster-1",
						},
					},
					DMMasterSpec: &apiv1.DMSpec{SpecCode: "master-spec"},
					DMWorkerSpec: &apiv1.DMSpec{SpecCode: "worker-spec"},
				},
			}

			masterNodeLabels := baseTask.GetMasterNodeLabels()
			g.Expect(masterNodeLabels).To(gomega.Equal(tt.expectedMasterLabel))

			workerNodeLabels := baseTask.GetWorkerNodeLabels()
			g.Expect(workerNodeLabels).To(gomega.Equal(tt.expectedWorkerLabel))
		})
	}
}