package eks

import (
	"testing"

	. "github.com/onsi/gomega"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"

	"github.com/tidbcloud/infra-provider/pkg/consts"
)

func TestAppendZonesLabelForNodegroup(t *testing.T) {
	cases := map[string]struct {
		ng     *cloudv1.NodegroupSpec
		result map[string]string
	}{
		"no zones": {
			ng: &cloudv1.NodegroupSpec{
				Zones:  nil,
				Labels: map[string]string{"k1": "v1"},
			},
			result: map[string]string{"k1": "v1"},
		},
		"one zone, nil label before": {
			ng: &cloudv1.NodegroupSpec{
				Zones: []string{"us-west-2a"},
			},
			result: map[string]string{consts.NodeGroupZonesLabelKey: "us-west-2a"},
		},
		"one zone": {
			ng: &cloudv1.NodegroupSpec{
				Zones:  []string{"us-west-2a"},
				Labels: map[string]string{"k1": "v1"},
			},
			result: map[string]string{"k1": "v1", consts.NodeGroupZonesLabelKey: "us-west-2a"},
		},
		"two zone": {
			ng: &cloudv1.NodegroupSpec{
				Zones:  []string{"us-west-2a", "us-west-2b"},
				Labels: map[string]string{"k1": "v1"},
			},
			result: map[string]string{"k1": "v1", consts.NodeGroupZonesLabelKey: "us-west-2a,us-west-2b"},
		},
		"two zone, different order": {
			ng: &cloudv1.NodegroupSpec{
				Zones:  []string{"us-west-2b", "us-west-2a"},
				Labels: map[string]string{"k1": "v1"},
			},
			result: map[string]string{"k1": "v1", consts.NodeGroupZonesLabelKey: "us-west-2a,us-west-2b"},
		},
		"one zone, overwrite, should not happened in real case": {
			ng: &cloudv1.NodegroupSpec{
				Zones:  []string{"us-west-2a"},
				Labels: map[string]string{consts.NodeGroupZonesLabelKey: "v1"},
			},
			result: map[string]string{consts.NodeGroupZonesLabelKey: "us-west-2a"},
		},
	}

	for name, cs := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			appendZonesLabelForNodegroup(cs.ng)
			g.Expect(cs.ng.Labels).To(Equal(cs.result))
		})
	}
}