package common

import (
	"testing"

	"github.com/stretchr/testify/require"

	clustertask "github.com/tidbcloud/eks-provider/controllers/cluster/task"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/namer"
)

func TestPrivateEndpointSourceRangesFollowAddressType(t *testing.T) {
	task := NewTiDBGroupService(&clustertask.Context{Cluster: &cloudv1.Cluster{}}, nil)
	group := cloudv1.ExtraGroupSpec{Name: "extra"}

	privateEndpoint := task.newPrivateEndpoint(&group, &cloudv1.ExtraGroupStatus{}, namer.NewNamer(), false)
	require.Equal(t, []string{"0.0.0.0/0"}, privateEndpoint.loadBalancerSourceRanges)

	group.SqlEndpoints.PrivateLinkNLBIPAddressType = cloudv1.PrivateLinkIPAddressTypeIPv4
	privateEndpoint = task.newPrivateEndpoint(&group, &cloudv1.ExtraGroupStatus{}, namer.NewNamer(), false)
	require.Equal(t, []string{"0.0.0.0/0"}, privateEndpoint.loadBalancerSourceRanges)

	group.SqlEndpoints.PrivateLinkNLBIPAddressType = cloudv1.PrivateLinkIPAddressTypeDualstack
	privateEndpoint = task.newPrivateEndpoint(&group, &cloudv1.ExtraGroupStatus{}, namer.NewNamer(), false)
	require.Equal(t, []string{"0.0.0.0/0", "::/0"}, privateEndpoint.loadBalancerSourceRanges)
}