package slprober

import (
	"context"
	"errors"
	"sync"
	"testing"
	"time"

	opv2 "github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes/fake"
	ctlrfake "sigs.k8s.io/controller-runtime/pkg/client/fake"

	"github.com/tidbcloud/sli-reporter/slprober/handshake"
)

// MockHandshakeClient is a mock implementation of handshake.Interface
type MockHandshakeClient struct {
	mock.Mock
}

func (m *MockHandshakeClient) TryConnect(ctx context.Context, addr string) (*handshake.HandshakePacket, error) {
	args := m.Called(ctx, addr)
	packet, _ := args.Get(0).(*handshake.HandshakePacket)
	return packet, args.Error(1)
}

func TestProbeDedicatedCluster_SkipProbe(t *testing.T) {
	mockKubeCli := fake.NewSimpleClientset()
	mockHandshakeCli := &MockHandshakeClient{}
	prober := &ServiceLevelProber{
		kubeCli: mockKubeCli,
	}

	tc := &v1alpha1.TidbCluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "db",
			Namespace: "tidb123",
		},
		Spec: v1alpha1.TidbClusterSpec{
			TiDB: &v1alpha1.TiDBSpec{
				Replicas: 1,
			},
		},
		Status: v1alpha1.TidbClusterStatus{
			TiDB: v1alpha1.TiDBStatus{
				Phase: v1alpha1.SuspendPhase,
			},
		},
	}
	tcList := []*v1alpha1.TidbCluster{tc}

	ctx := context.Background()
	var wg sync.WaitGroup
	resultCh := make(chan clusterResult, 1)

	wg.Add(1)
	go prober.probeDedicatedCluster(ctx, tcList, &wg, resultCh, mockHandshakeCli)
	wg.Wait()
	close(resultCh)

	result := <-resultCh
	assert.Equal(t, ServiceStatusSkippedBySuspend, result.status)
	assert.Equal(t, getClusterID(tc), result.clusterID)
}

func TestProbeDedicatedCluster(t *testing.T) {
	tests := []struct {
		name            string
		tcList          []*v1alpha1.TidbCluster
		svcList         []v1.Service
		handshakeErrors map[string]error
		expectedStatus  ServiceStatus
	}{
		{
			name: "one tc, alive",
			tcList: []*v1alpha1.TidbCluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db",
						Namespace: "tidb123",
					},
					Spec: v1alpha1.TidbClusterSpec{
						TiDB: &v1alpha1.TiDBSpec{
							Replicas: 1,
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-tidb",
						Namespace: "tidb123",
						Labels: map[string]string{
							"app.kubernetes.io/used-by":    "end-user",
							"app.kubernetes.io/name":       "tidb-cluster",
							"app.kubernetes.io/managed-by": "tidb-operator",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": nil,
			},
			expectedStatus: ServiceStatusDetectedAsAlive,
		},
		{
			name: "one tc, dead",
			tcList: []*v1alpha1.TidbCluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db",
						Namespace: "tidb123",
					},
					Spec: v1alpha1.TidbClusterSpec{
						TiDB: &v1alpha1.TiDBSpec{
							Replicas: 1,
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-tidb",
						Namespace: "tidb123",
						Labels: map[string]string{
							"app.kubernetes.io/used-by":    "end-user",
							"app.kubernetes.io/name":       "tidb-cluster",
							"app.kubernetes.io/managed-by": "tidb-operator",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": errors.New("connection failed"),
			},
			expectedStatus: ServiceStatusDetectedAsDead,
		},
		{
			name: "one tc, timeout",
			tcList: []*v1alpha1.TidbCluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db",
						Namespace: "tidb123",
					},
					Spec: v1alpha1.TidbClusterSpec{
						TiDB: &v1alpha1.TiDBSpec{
							Replicas: 1,
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-tidb",
						Namespace: "tidb123",
						Labels: map[string]string{
							"app.kubernetes.io/used-by":    "end-user",
							"app.kubernetes.io/name":       "tidb-cluster",
							"app.kubernetes.io/managed-by": "tidb-operator",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": context.DeadlineExceeded,
			},
			expectedStatus: ServiceStatusDetectedAsDead,
		},
		{
			name: "multi tc, all alive",
			tcList: []*v1alpha1.TidbCluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db",
						Namespace: "tidb123",
					},
					Spec: v1alpha1.TidbClusterSpec{
						TiDB: &v1alpha1.TiDBSpec{
							Replicas: 1,
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-222",
						Namespace: "tidb123",
					},
					Spec: v1alpha1.TidbClusterSpec{
						TiDB: &v1alpha1.TiDBSpec{
							Replicas: 1,
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-tidb",
						Namespace: "tidb123",
						Labels: map[string]string{
							"app.kubernetes.io/used-by":    "end-user",
							"app.kubernetes.io/name":       "tidb-cluster",
							"app.kubernetes.io/managed-by": "tidb-operator",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-222-tidb",
						Namespace: "tidb123",
						Labels: map[string]string{
							"app.kubernetes.io/used-by":    "end-user",
							"app.kubernetes.io/name":       "tidb-cluster",
							"app.kubernetes.io/managed-by": "tidb-operator",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "5.6.7.8",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": nil,
				"5.6.7.8:4000": nil,
			},
			expectedStatus: ServiceStatusDetectedAsAlive,
		},
		{
			name: "multi tc, one dead",
			tcList: []*v1alpha1.TidbCluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db",
						Namespace: "tidb123",
					},
					Spec: v1alpha1.TidbClusterSpec{
						TiDB: &v1alpha1.TiDBSpec{
							Replicas: 1,
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-222",
						Namespace: "tidb123",
					},
					Spec: v1alpha1.TidbClusterSpec{
						TiDB: &v1alpha1.TiDBSpec{
							Replicas: 1,
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-tidb",
						Namespace: "tidb123",
						Labels: map[string]string{
							"app.kubernetes.io/used-by":    "end-user",
							"app.kubernetes.io/name":       "tidb-cluster",
							"app.kubernetes.io/managed-by": "tidb-operator",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-222-tidb",
						Namespace: "tidb123",
						Labels: map[string]string{
							"app.kubernetes.io/used-by":    "end-user",
							"app.kubernetes.io/name":       "tidb-cluster",
							"app.kubernetes.io/managed-by": "tidb-operator",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "5.6.7.8",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": nil,
				"5.6.7.8:4000": errors.New("connection failed"),
			},
			expectedStatus: ServiceStatusDetectedAsAlive,
		},
		{
			name: "multi tc, all dead",
			tcList: []*v1alpha1.TidbCluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db",
						Namespace: "tidb123",
					},
					Spec: v1alpha1.TidbClusterSpec{
						TiDB: &v1alpha1.TiDBSpec{
							Replicas: 1,
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-222",
						Namespace: "tidb123",
					},
					Spec: v1alpha1.TidbClusterSpec{
						TiDB: &v1alpha1.TiDBSpec{
							Replicas: 1,
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-tidb",
						Namespace: "tidb123",
						Labels: map[string]string{
							"app.kubernetes.io/used-by":    "end-user",
							"app.kubernetes.io/name":       "tidb-cluster",
							"app.kubernetes.io/managed-by": "tidb-operator",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-222-tidb",
						Namespace: "tidb123",
						Labels: map[string]string{
							"app.kubernetes.io/used-by":    "end-user",
							"app.kubernetes.io/name":       "tidb-cluster",
							"app.kubernetes.io/managed-by": "tidb-operator",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "5.6.7.8",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": errors.New("connection failed"),
				"5.6.7.8:4000": errors.New("connection failed"),
			},
			expectedStatus: ServiceStatusDetectedAsDead,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mockKubeCli := fake.NewSimpleClientset()
			for _, svc := range tt.svcList {
				_, _ = mockKubeCli.CoreV1().
					Services(svc.Namespace).
					Create(context.Background(), &svc, metav1.CreateOptions{})
			}
			// shouldTimeout := false
			mockHandshakeCli := &MockHandshakeClient{}
			for addr, err := range tt.handshakeErrors {
				if errors.Is(err, context.DeadlineExceeded) {
					// shouldTimeout = true
					mockHandshakeCli.On("TryConnect", mock.Anything, addr).After(1*time.Second).Return(nil, err)
				} else {
					mockHandshakeCli.On("TryConnect", mock.Anything, addr).Return(nil, err)
				}
			}

			ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
			defer cancel()
			prober := &ServiceLevelProber{
				kubeCli: mockKubeCli,
			}
			var wg sync.WaitGroup
			resultCh := make(chan clusterResult, 1)

			wg.Add(1)
			go prober.probeDedicatedCluster(ctx, tt.tcList, &wg, resultCh, mockHandshakeCli)
			wg.Wait()
			close(resultCh)

			result := <-resultCh
			assert.Equal(t, tt.expectedStatus, result.status)
			assert.Equal(t, getClusterID(tt.tcList[0]), result.clusterID)

			time.Sleep(200 * time.Millisecond) // Wait for goroutines to finish
			mockHandshakeCli.AssertExpectations(t)
		})
	}
}

func TestProbeNextGenCluster_SkipProbe(t *testing.T) {
	mockKubeCli := fake.NewSimpleClientset()
	mockHandshakeCli := &MockHandshakeClient{}
	scheme := runtime.NewScheme()
	mockCli := ctlrfake.NewClientBuilder().WithScheme(scheme).Build()

	prober := &ServiceLevelProber{
		kubeCli: mockKubeCli,
		cli:     mockCli,
	}

	tidbGroupList := []*opv2.TiDBGroup{
		{
			ObjectMeta: metav1.ObjectMeta{
				Labels: map[string]string{
					LabelCluster: "12345",
				},
				Name:      "db-1111",
				Namespace: "tidb12345",
			},
			Status: opv2.TiDBGroupStatus{
				GroupStatus: opv2.GroupStatus{
					ReadyReplicas: 0,
				},
			},
		},
		{
			ObjectMeta: metav1.ObjectMeta{
				Labels: map[string]string{
					LabelCluster: "12345",
				},
				Name:      "db-2222",
				Namespace: "tidb12345",
			},
			Status: opv2.TiDBGroupStatus{
				GroupStatus: opv2.GroupStatus{
					ReadyReplicas: 0,
				},
			},
		},
	}

	ctx := context.Background()
	var wg sync.WaitGroup
	resultCh := make(chan clusterResult, 1)

	wg.Add(1)
	go prober.probeNextGenCluster(ctx, tidbGroupList, &wg, resultCh, mockHandshakeCli)
	wg.Wait()
	close(resultCh)

	result := <-resultCh
	assert.Equal(t, ServiceStatusSkippedByNotReady, result.status)
}

func TestServiceLevelProber_probeNextGenCluster(t *testing.T) {
	scheme := runtime.NewScheme()
	err := opv2.Install(scheme)
	if err != nil {
		t.Fatalf("failed to add to schema: %v", err.Error())
	}

	tests := []struct {
		name             string
		tidbGroupList    []*opv2.TiDBGroup
		tiproxyGroupList []*opv2.TiProxyGroup
		svcList          []v1.Service
		handshakeErrors  map[string]error
		expectedStatus   ServiceStatus
	}{
		{
			name: "1 TiDBGroup for dedicated, all alive, result: alive",
			tidbGroupList: []*opv2.TiDBGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 2,
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-tikv",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "2.3.4.5",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 8000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": nil,
			},
			expectedStatus: ServiceStatusDetectedAsAlive,
		},
		{
			name: "2 TiDBGroup, 1 TiProxyGroup, all alive, result: alive",
			tidbGroupList: []*opv2.TiDBGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-1111",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-2222",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
			},
			tiproxyGroupList: []*opv2.TiProxyGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelCluster: "12345",
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-1111-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-1111",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-2222-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-2222",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "2.3.4.5",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333-tiproxy",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-3333",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "3.4.5.6",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": nil,
				"2.3.4.5:4000": nil,
				"3.4.5.6:4000": nil,
			},
			expectedStatus: ServiceStatusDetectedAsAlive,
		},
		{
			name: "2 TiDBGroup, 1 TiProxyGroup, 1 TiDBGroup dead, others alive, result: alive",
			tidbGroupList: []*opv2.TiDBGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-1111",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-2222",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
			},
			tiproxyGroupList: []*opv2.TiProxyGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelCluster: "12345",
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-1111-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-1111",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-2222-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-2222",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "2.3.4.5",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333-tiproxy",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-3333",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "3.4.5.6",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": errors.New("connection failed"),
				"2.3.4.5:4000": nil,
				"3.4.5.6:4000": nil,
			},
			expectedStatus: ServiceStatusDetectedAsAlive,
		},
		{
			name: "2 TiDBGroup, 1 TiProxyGroup, 1 TiDBGroup timeout, others alive, result: alive",
			tidbGroupList: []*opv2.TiDBGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-1111",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-2222",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
			},
			tiproxyGroupList: []*opv2.TiProxyGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelCluster: "12345",
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-1111-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-1111",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-2222-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-2222",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "2.3.4.5",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333-tiproxy",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-3333",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "3.4.5.6",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": nil,
				"2.3.4.5:4000": context.DeadlineExceeded,
				"3.4.5.6:4000": nil,
			},
			expectedStatus: ServiceStatusDetectedAsAlive,
		},
		{
			name: "2 TiDBGroup, 1 TiProxyGroup, 2 TiDBGroup dead, others alive, result: dead",
			tidbGroupList: []*opv2.TiDBGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-1111",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-2222",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
			},
			tiproxyGroupList: []*opv2.TiProxyGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelCluster: "12345",
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-1111-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-1111",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-2222-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-2222",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "2.3.4.5",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333-tiproxy",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-3333",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "3.4.5.6",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": errors.New("connection failed"),
				"2.3.4.5:4000": context.DeadlineExceeded,
				"3.4.5.6:4000": nil,
			},
			expectedStatus: ServiceStatusDetectedAsDead,
		},
		{
			name: "2 TiDBGroup, 1 TiProxyGroup, 1 TiProxyGroup dead, others alive, result: dead",
			tidbGroupList: []*opv2.TiDBGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-1111",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-2222",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
			},
			tiproxyGroupList: []*opv2.TiProxyGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelCluster: "12345",
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-1111-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-1111",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-2222-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-2222",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "2.3.4.5",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333-tiproxy",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-3333",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "3.4.5.6",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": nil,
				"2.3.4.5:4000": nil,
				"3.4.5.6:4000": errors.New("connection failed"),
			},
			expectedStatus: ServiceStatusDetectedAsDead,
		},
		{
			name: "2 TiDBGroup, 1 TiProxyGroup, all dead, result: dead",
			tidbGroupList: []*opv2.TiDBGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-1111",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-2222",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
			},
			tiproxyGroupList: []*opv2.TiProxyGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelCluster: "12345",
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-1111-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-1111",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-2222-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-2222",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "2.3.4.5",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333-tiproxy",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-3333",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "3.4.5.6",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": errors.New("connection failed"),
				"2.3.4.5:4000": context.DeadlineExceeded,
				"3.4.5.6:4000": context.DeadlineExceeded,
			},
			expectedStatus: ServiceStatusDetectedAsDead,
		},
		{
			name: "2 TiDBGroup, 1 TiProxyGroup, 1 TiDBGroup doesn't have a svc, others alive, result: alive",
			tidbGroupList: []*opv2.TiDBGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-1111",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-2222",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
			},
			tiproxyGroupList: []*opv2.TiProxyGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelCluster: "12345",
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-1111-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-1111",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-2233-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-2233",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "2.3.4.5",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333-tiproxy",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-3333",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "3.4.5.6",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": nil,
				// "2.3.4.5:4000": nil,
				"3.4.5.6:4000": nil,
			},
			expectedStatus: ServiceStatusDetectedAsAlive,
		},
		{
			name: "2 TiDBGroup, 1 TiProxyGroup, 1 TiProxyGroup doesn't have a svc, others alive, result: alive",
			tidbGroupList: []*opv2.TiDBGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-1111",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							LabelCluster: "12345",
						},
						Name:      "db-2222",
						Namespace: "tidb12345",
					},
					Status: opv2.TiDBGroupStatus{
						GroupStatus: opv2.GroupStatus{
							ReadyReplicas: 1,
						},
					},
				},
			},
			tiproxyGroupList: []*opv2.TiProxyGroup{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3333",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelCluster: "12345",
						},
					},
				},
			},
			svcList: []v1.Service{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-1111-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-1111",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "1.2.3.4",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-2222-tidb",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-2222",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "2.3.4.5",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "db-3344-tiproxy",
						Namespace: "tidb12345",
						Labels: map[string]string{
							LabelGroup: "db-3344",
						},
					},
					Spec: v1.ServiceSpec{
						ClusterIP: "3.4.5.6",
						Ports: []v1.ServicePort{
							{Name: "mysql-client", Port: 4000},
						},
					},
				},
			},
			handshakeErrors: map[string]error{
				"1.2.3.4:4000": nil,
				"2.3.4.5:4000": nil,
				// "3.4.5.6:4000": nil,
			},
			expectedStatus: ServiceStatusDetectedAsAlive,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// mock to create svc
			mockKubeCli := fake.NewSimpleClientset()
			for _, svc := range tt.svcList {
				_, _ = mockKubeCli.CoreV1().
					Services(svc.Namespace).
					Create(context.Background(), &svc, metav1.CreateOptions{})
			}
			mockCli := ctlrfake.NewClientBuilder().WithScheme(scheme).Build()

			// mock to create TiProxyGroup
			for _, tg := range tt.tiproxyGroupList {
				err := mockCli.Create(context.Background(), tg)
				assert.NoError(t, err, "failed to create TiProxyGroup")
			}

			// shouldTimeout := false
			mockHandshakeCli := &MockHandshakeClient{}
			for addr, err := range tt.handshakeErrors {
				if errors.Is(err, context.DeadlineExceeded) {
					// shouldTimeout = true
					mockHandshakeCli.On("TryConnect", mock.Anything, addr).After(1*time.Second).Return(nil, err)
				} else {
					mockHandshakeCli.On("TryConnect", mock.Anything, addr).Return(nil, err)
				}
			}

			ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
			defer cancel()
			prober := &ServiceLevelProber{
				kubeCli: mockKubeCli,
				cli:     mockCli,
			}
			var wg sync.WaitGroup
			resultCh := make(chan clusterResult, 1)

			wg.Add(1)
			go prober.probeNextGenCluster(ctx, tt.tidbGroupList, &wg, resultCh, mockHandshakeCli)
			wg.Wait()
			close(resultCh)

			result := <-resultCh
			assert.Equal(t, tt.expectedStatus, result.status)
			if _, ok := tt.tidbGroupList[0].Labels[LabelCluster]; ok {
				assert.Equal(t, tt.tidbGroupList[0].Labels[LabelCluster], result.clusterID)
			} else {
				assert.Equal(t, tt.tidbGroupList[0].Namespace[4:], result.clusterID)
			}

			time.Sleep(200 * time.Millisecond) // Wait for goroutines to finish
			mockHandshakeCli.AssertExpectations(t)
		})
	}
}