package slprober

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

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
)

func TestProbeTiDBService_Success(t *testing.T) {
	mockHandshake := new(MockHandshakeClient)
	mockHandshake.On("TryConnect", mock.Anything, "127.0.0.1:4000").Return(nil, nil)

	resultCh := make(chan serviceResult, 1)
	ctx := context.Background()

	go probeTiDBService(ctx, "127.0.0.1:4000", "tidb", resultCh, mockHandshake)

	select {
	case result := <-resultCh:
		assert.True(t, result.success)
		assert.Equal(t, "127.0.0.1:4000", result.addr)
		assert.Equal(t, "tidb", result.serviceName)
	case <-time.After(1 * time.Second):
		t.Fatal("Test timed out")
	}

	mockHandshake.AssertExpectations(t)
}

func TestProbeTiDBService_Failure(t *testing.T) {
	mockHandshake := new(MockHandshakeClient)
	mockHandshake.On("TryConnect", mock.Anything, "127.0.0.1:4000").Return(nil, errors.New("connection failed"))

	resultCh := make(chan serviceResult, 1)
	ctx := context.Background()

	go probeTiDBService(ctx, "127.0.0.1:4000", "tidb", resultCh, mockHandshake)

	select {
	case result := <-resultCh:
		assert.False(t, result.success)
		assert.Equal(t, "127.0.0.1:4000", result.addr)
		assert.Equal(t, "tidb", result.serviceName)
	case <-time.After(1 * time.Second):
		t.Fatal("Test timed out")
	}

	mockHandshake.AssertExpectations(t)
}

func TestProbeTiDBService_ContextDone(t *testing.T) {
	mockHandshake := new(MockHandshakeClient)
	mockHandshake.On("TryConnect", mock.Anything, "127.0.0.1:4000").After(1*time.Second).Return(nil, nil)

	resultCh := make(chan serviceResult, 1)
	ctx, cancel := context.WithCancel(context.Background())
	cancel() // Cancel the context immediately

	go probeTiDBService(ctx, "127.0.0.1:4000", "tidb", resultCh, mockHandshake)

	select {
	case <-resultCh:
		t.Fatal("Expected no result due to context cancellation")
	case <-time.After(500 * time.Millisecond):
		// Test passes if no result is sent
	}

	mockHandshake.AssertExpectations(t)
}

func TestProbeTiDBService_Timeout(t *testing.T) {
	mockHandshake := new(MockHandshakeClient)
	mockHandshake.On("TryConnect", mock.Anything, "127.0.0.1:4000").After(1*time.Second).Return(nil, nil)

	resultCh := make(chan serviceResult, 1)
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
	defer cancel()

	go probeTiDBService(ctx, "127.0.0.1:4000", "tidb", resultCh, mockHandshake)
	select {
	case <-resultCh:
		t.Fatal("Expected no result due to context cancellation")
	case <-time.After(200 * time.Millisecond):
		// Test passes if no result is sent
	}

	mockHandshake.AssertExpectations(t)
}