package entity

import (
	"testing"

	"github.com/stretchr/testify/require"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
)

func TestBuildTiProxyGracefulRestartLoadBalancerAnnotations(t *testing.T) {
	tests := []struct {
		name     string
		provider apiv1.ProviderType
		expected map[string]string
	}{
		{
			name:     "aws",
			provider: apiv1.ProviderType_AWS,
			expected: map[string]string{
				AWSLBHealthcheckPathKey:     TiProxyHealthcheckPath,
				AWSLBHealthcheckPortKey:     TiProxyHealthcheckPort,
				AWSLBHealthcheckProtocolKey: "http",
			},
		},
		{
			name:     "aliyun",
			provider: apiv1.ProviderType_ALICLOUD,
			expected: map[string]string{
				AliCloudLBHealthcheckFlagKey:           "on",
				AliCloudLBHealthcheckTypeKey:           TiProxyAliCloudHealthcheckProtocol,
				AliCloudLBHealthcheckURIKey:            TiProxyHealthcheckPath,
				AliCloudLBHealthcheckDomainKey:         "$SERVER_IP",
				AliCloudLBHealthcheckConnectPortKey:    TiProxyHealthcheckPort,
				AliCloudLBHealthcheckConnectTimeoutKey: "5",
				AliCloudLBHealthyThresholdKey:          "2",
				AliCloudLBUnhealthyThresholdKey:        "2",
				AliCloudLBHealthcheckIntervalKey:       "5",
				AliCloudLBHealthcheckMethodKey:         "GET",
			},
		},
		{
			name:     "azure",
			provider: apiv1.ProviderType_AZURE,
			expected: map[string]string{
				"service.beta.kubernetes.io/port_4000_health-probe_protocol":     "http",
				"service.beta.kubernetes.io/port_4000_health-probe_port":         TiProxyHealthcheckPort,
				"service.beta.kubernetes.io/port_4000_health-probe_request-path": TiProxyHealthcheckPath,
				"service.beta.kubernetes.io/port_4000_health-probe_interval":     "5",
				"service.beta.kubernetes.io/port_4000_health-probe_num-of-probe": "2",
			},
		},
		{
			name:     "unsupported provider",
			provider: apiv1.ProviderType_GCP,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			sharedPool := &SharedPool{
				SharedPool: apiv1.SharedPool{Provider: &apiv1.Provider{Type: tt.provider}},
			}

			require.Equal(t, tt.expected, BuildTiProxyGracefulRestartLoadBalancerAnnotations(sharedPool, 4000))
		})
	}
}

func TestBuildTiProxyGracefulRestartLoadBalancerAnnotationsAzureRequiresServicePort(t *testing.T) {
	sharedPool := &SharedPool{
		SharedPool: apiv1.SharedPool{Provider: &apiv1.Provider{Type: apiv1.ProviderType_AZURE}},
	}

	require.Nil(t, BuildTiProxyGracefulRestartLoadBalancerAnnotations(sharedPool))
	require.Nil(t, BuildTiProxyGracefulRestartLoadBalancerAnnotations(sharedPool, 0))
}